task77_installing_a_modem/README.md

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

Single-task fine-tune for the [BEHAVIOR-1K](https://behavior.stanford.edu/) 2026 challenge.
Part of [`JackLiu0406/b1k-checkpoints`](https://huggingface.co/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`](https://huggingface.co/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:

```python
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

```bash
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

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

### 3. Serve the policy

```bash
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

```python
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.