# RX1 v2 — learned 40-DoF action model

v2 replaces the scripted-IK controller with a **learned action model** that
drives the robot. All v1 files are untouched (`scripts/rx1_brain_server.py`,
`models/policy.py`, …); v2 lives in new files alongside them.

## What drives the robot now

```
camera ─▶ Qwen2.5-VL-3B  (perception + high-level plan: which cube, what target)
              │            models/vlm_planner.py   (GPU, optional; keyword fallback)
              ▼
proprio ─▶ Learned action model  (40-DoF arms + grasp, goal-conditioned)
              │   models/act_model.py · models/policy_learned.py
              ▼
        joint targets ─▶ MuJoCo (RX1Env)
```

- **Action model** (`models/act_model.py`): goal-conditioned, action-chunking MLP.
  `obs = [proprio(80), pick_pos0(3), place_pos(3)] → trajectory (T,15)` where each
  step is `[7 right-arm, 7 left-arm, 1 grasp]`. Torso/head are held compliant and
  the 18 finger joints follow the grasp bit (the RX1 arm can't reach the table
  from a rigid torso — see the design note in `models/expert.py`).
- It is queried **once per pick-place step** and the trajectory is played
  **open-loop** (`models/policy_learned.py`), which sidesteps behaviour-cloning
  covariate shift. One model handles pick / place / sort / stack — only the place
  target differs.

## Train it (your GPU)

```bash
pip install -r requirements_v2.txt

# 1. collect demonstrations from the privileged IK teacher (CPU, ~minutes)
#    free-space pick/place/sort:
python scripts/collect_demos.py --n 800 --out data/demos.npz
#    add stacking demos (base cube under the target) to learn cube-on-cube:
python scripts/collect_demos.py --n 800 --out data/demos_stack.npz --base-frac 0.5

# 2. behaviour-clone the action model (GPU recommended)
python scripts/train_act.py --demos data/demos.npz --out models/ckpt/act.pt \
    --epochs 400 --chunk 216

# 3. run the learned brain server
python scripts/rx1_brain_server_v2.py --render          # MuJoCo viewer
#    headless:  python scripts/rx1_brain_server_v2.py
curl -X POST http://localhost:8788/task \
  -H 'Content-Type: application/json' -d '{"instruction":"sort the cubes"}'
```

Tasks understood (free-form NL, any of red/blue/green/yellow):
`pick/place/move <colour> [in box | on <colour>]`, `sort`/`tidy`/`clear`,
`stack [the cubes | <a> on <b>]`, plus `stop`.

## Status (honest)

Verified end-to-end on a **CPU-trained** model (a few hundred episodes):

- **Pick-and-place and sorting into the box work** — the learned model places a
  cube in the box at ~0.02–0.05 m through the full brain→policy→MuJoCo loop.
- **Tight cube-on-cube stacking is precision-limited.** The CPU model places at
  ~0.05 m; a 0.044 m cube needs ~0.02 m to stay on another cube, and the base
  cube's contact deflects a slightly-off placement. Stacking improves with (a)
  the `--base-frac` stacking demos above, (b) far more episodes, and (c) longer
  training on your 3060 — the pipeline supports all three. This is a data/compute
  ceiling, not a code limit.

The high-level planner falls back to a keyword parser if Qwen2.5-VL / a GPU is
unavailable, and `policy_learned.py` falls back to the v1 IK policy if no trained
checkpoint is present, so the server always runs.