task77_installing_a_modem/README.md

Task 77 — installing_a_modem (PiBehavior / pi0.5, no DA3)

Single-task fine-tune for the BEHAVIOR-1K 2026 challenge. Part of JackLiu0406/b1k-checkpoints.

This is one checkpoint line in the repo; the DA3 K/V-split spatial models are documented in the root README.

---

Quick facts

architecture	PiBehavior / π0.5 — gemma_2b VLM + gemma_300m action expert
parameters	3,826.8 M
framework	JAX / Flax NNX, Orbax checkpoint
initialised from	IliaLarchenko/behavior_50t_checkpoint (50-task meta-trained)
fine-tuned on	task 77 = installing_a_modem, 200 episodes, ~476 k samples
checkpoint step	29,000 (of a 30,000-step schedule — see Caveats)
task-embedding space	100 tasks (expanded from the upstream 50)
DA3	not used — this is the plain π0.5 path
precision	params stored as saved by Orbax; model runs bf16 compute

Final training metrics (step 29,275, last logged)

metric	value
action_loss	0.0081
total_loss	0.0097
fast_accuracy	0.9932
subtask_accuracy	1.0000

---

Repo layout

task77_installing_a_modem/
└── 29000/
    ├── params/                 # Orbax OCDBT param tree (~11.8 GiB) — inference weights
    ├── assets/
    │   └── IliaLarchenko/
    │       └── behavior_224_rgb/
    │           ├── norm_stats.json     # state/action normalisation
    │           └── fast_tokenizer/     # FAST action tokenizer (auxiliary head)
    └── _CHECKPOINT_METADATA

train_state/ is deliberately not included. That directory is 31 GB of Adam optimiser moments, needed only to resume training. Everything required for inference and evaluation is in params/.

---

Input contract

This is the part most people get wrong, so it is spelled out exactly.

⚠️ There is no text prompt

PiBehavior does not take a language instruction. It is conditioned by a task embedding and a stage (subtask) embedding, both looked up by integer index. Any prompt string you pass is ignored — serve_b1k.py sets a placeholder string purely for logging.

Observation dict

key	shape	dtype	notes
images["base_0_rgb"]	[B, 224, 224, 3]	float32	head/base camera
images["left_wrist_0_rgb"]	[B, 224, 224, 3]	float32	left wrist camera
images["right_wrist_0_rgb"]	[B, 224, 224, 3]	float32	right wrist camera
image_masks[<same three keys>]	[B]	bool	one flag per image, not per pixel — False marks a missing camera
state	[B, 32]	float32	proprioception, normalised with norm_stats.json
tokenized_prompt	[B, 2]	int32	[task_id, subtask_state]
tokenized_prompt_mask	[B, 2]	bool	normally [True, True]

fast_tokens / fast_token_mask appear only when the FAST auxiliary loss is enabled; they are a training-time auxiliary target and are not needed at inference.

Image resolution is openpi.models.model.IMAGE_RESOLUTION == (224, 224).

Output

key	shape	dtype
actions	[B, 30, 32]	float32

That is an action chunk: horizon 30, action dimension 32. Denormalise with the actions entry of norm_stats.json. The 32 dims cover base velocity (x, y, z), left arm (7) + left gripper, right arm (7) + right gripper, and trunk (4), with the remainder padding to 32.

---

Task and stage indexing — important

The upstream release has a 50-task embedding table. installing_a_modem is a 2026-only activity and is not in those 50, so this model was trained with the table expanded to 100 tasks; rows 50–99 were freshly initialised from normal(0, 1/√features) and rows 0–49 were carried over unchanged.

You must load this checkpoint with a 100-task config, e.g. B1K_TASK_SPACE=100. Loading it into a 50-task model will fail on a shape mismatch.

Stage counts follow the verified upstream rule:

num_stages = clip(ceil(avg_episode_length / 900), 5, 15)

For this model:

	value
task_id	77
num_stages	5
stage-embedding rows	888 – 892

subtask_state at inference is the current stage index in 0 … num_stages-1 (so 0–4 here). The absolute embedding row is TASK_STAGE_OFFSETS[task_id] + subtask_state; the offsets are a cumulative sum over TASK_NUM_STAGES (100 entries, summing to 1120).

---

Running it

The evaluation path is a websocket policy server that the OmniGibson BEHAVIOR-1K evaluator connects to.

1. Get the code

git clone https://github.com/IliaLarchenko/behavior-1k-solution
cd behavior-1k-solution
git submodule update --init --recursive
uv sync

This checkpoint was produced by a fork of the above with the task table widened to 100. If you use stock upstream, you must apply the same widening or the parameter shapes will not match.

2. Download the checkpoint

huggingface-cli download <REPO_ID> \
  --include "task77_installing_a_modem/29000/*" \
  --local-dir ./ckpts

3. Serve the policy

export B1K_TASK_SPACE=100                       # 100-task embedding table
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

uv run scripts/serve_b1k.py \
  --policy.config pi_behavior_b1k_fast \
  --policy.dir   ./ckpts/task77_installing_a_modem/29000 \
  --task_id 77 \
  --port 8000

--task_id 77 is what selects the task embedding. Omit it only if your observations already carry a task_index / task_id field, in which case it is read per-observation.

assets_base_dir must resolve so that norm_stats.json and fast_tokenizer/ are found under <assets>/IliaLarchenko/behavior_224_rgb/. The bundled assets/ directory already has that layout — point the config at it if the default path does not exist on your machine.

Eval-side execution parameters

serve_b1k.py wraps the raw policy in B1KPolicyWrapper, which does chunk scheduling and a few eval heuristics. Defaults:

arg	default	meaning
--actions_to_execute	26	actions consumed per inference
--actions_to_keep	4	overlap retained between chunks
--execute_in_n_steps	20	env steps per executed chunk
--history_len	3	observation history for stage voting
--votes_to_promote	2	consecutive votes before advancing a stage
--num_steps	20	flow-matching sampling steps
--apply_eval_tricks	True	correction rules + gripper variation checks

Set --apply_eval_tricks False for a clean measurement of the policy itself.

Loading programmatically

from b1k.policies import policy_config as _policy_config
from b1k.training import config as _config

policy = _policy_config.create_trained_policy(
    _config.get_config("pi_behavior_b1k_fast"),
    "./ckpts/task77_installing_a_modem/29000",
    sample_kwargs={"num_steps": 20},
)

action_chunk = policy.infer({
    "images": {...},            # three [224,224,3] float32 arrays
    "image_masks": {...},       # three bools
    "state": state_32,          # [32] float32
    "task_id": 77,
    "subtask_state": current_stage,   # 0..4
})["actions"]                   # [30, 32]

---

Training recipe

init	IliaLarchenko/behavior_50t_checkpoint params, task/stage tables widened 50→100
GPUs	7 × H200 (one card on the node was faulty and excluded)
global batch	224 (32 per GPU)
FSDP	off — parameters replicated
steps	30,000 scheduled; stopped at 29,275, last checkpoint 29,000
LR schedule	ramp 8.75e-7 → peak 8.75e-5 over 2,000 warmup steps, cosine to 1.75e-6
LR scaling	linear rule (Goyal et al.), 1.75× over a BS-128 reference of 5e-5
optimiser	AdamW (openpi defaults)
new params	task_embeddings[77] and task_stage_embeddings[888:893] initialised from normal(0, 1/√features); everything else transferred

The FAST action tokenizer was not retrained. It was verified to transfer to the new tasks: 0.000 % alphabet overflow, no clipping, and round-trip MAE on the new task matching in-distribution controls.

---

Caveats — read before comparing numbers

1. This is step 29,000, not 30,000. Training was stopped 725 steps early to free the machine for hardware fault diagnosis. Those steps were at the LR floor (~1.75e-6) and loss was already flat, so the practical difference is very small — but it is not the full schedule.

2. Inference only. train_state/ is not included, so you cannot resume training from this upload.

3. Base-velocity coordinate frame. Trained on the 2026 challenge demos after the upstream fix (commit e6c9756) that moved base_qvel from the world frame to the robot frame. Checkpoints trained on the older world-frame data saw a different input distribution for state dims 0:3 — losses are not directly comparable across that boundary.

4. Single-task model. It has only been fine-tuned on task 77. The other 99 task embeddings are either inherited from the 50-task meta checkpoint (0–49) or still at random initialisation (50–99, except 77). Do not expect meaningful behaviour for any other task_id.

5. subtask_accuracy = 1.0000 is on training data — it is not a generalisation estimate.