feat(vla): add ACT policy for socket peg

2026-07-31 10:11:04 +08:00
parent acbd7c605a
commit ff7f4a1b03
8 changed files with 926 additions and 5 deletions
@@ -0,0 +1,49 @@
+# ACT Socket Peg Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Add a local ACT policy/model to RoboIMI and launch training on the socket peg dataset with three 224×224 camera views.
+
+**Architecture:** Implement a self-contained ACT agent and head that reuse the existing ResNet multiview backbone, dataset, training loop, normalization, and checkpointing. The ACT model uses a posterior transformer encoder for latent z and a transformer decoder with learned action queries for action chunks.
+
+**Tech Stack:** Python, PyTorch, Hydra/OmegaConf, unittest, HDF5 dataset via existing `SimpleRobotDataset`.
+
+---
+
+## File Structure
+- Create `roboimi/vla/models/heads/act.py`: local ACT model/head implementation, no imports from external ACT repository.
+- Create `roboimi/vla/agent_act.py`: VLA-compatible ACT agent wrapper with normalization, condition building, loss, and inference queues.
+- Create `roboimi/vla/conf/agent/act_resnet.yaml`: Hydra agent config for three-camera ACT with 224×224 images.
+- Create `tests/test_act_agent.py`: model/agent unit tests.
+- Modify no external ACT code and do not add vendored ACT files.
+
+## Tasks
+
+### Task 1: Add ACT model/head tests
+- [ ] Write tests in `tests/test_act_agent.py` that define a lightweight fake vision backbone emitting deterministic camera tokens.
+- [ ] Test `ACTAgent.compute_loss()` returns a scalar tensor and backpropagates through the head.
+- [ ] Test masked L1 ignores padded timesteps by comparing all-padded vs partially valid batches for finite loss behavior.
+- [ ] Test `ACTAgent.predict_action()` returns `(B,pred_horizon,action_dim)`.
+- [ ] Run `python -m unittest tests.test_act_agent -v` and confirm tests fail because `roboimi.vla.agent_act` does not exist.
+
+### Task 2: Implement local ACT head and ACT agent
+- [ ] Create `roboimi/vla/models/heads/act.py` with `ACTPolicyHead`, sinusoidal table helper, KL helper, and transformer layers using `batch_first=True` PyTorch modules.
+- [ ] Create `roboimi/vla/agent_act.py` with `ACTAgent` implementing existing training/inference API.
+- [ ] Reuse `NormalizationModule` and camera ordering checks from `VLAAgent` behavior.
+- [ ] Run `python -m unittest tests.test_act_agent -v` and fix until green.
+
+### Task 3: Add Hydra config and wiring tests
+- [ ] Add `roboimi/vla/conf/agent/act_resnet.yaml` using existing `resnet_diffusion` backbone with `output_tokens_per_camera=true` and `camera_names=${data.camera_names}`.
+- [ ] Extend `tests/test_act_agent.py` with a Hydra compose/instantiate test using reduced backbone/head sizes and `data.camera_names='[l_vis,r_vis,front]'`.
+- [ ] Run `python -m unittest tests.test_act_agent -v` and `python -m unittest tests.test_resnet_transformer_agent_wiring -v`.
+
+### Task 4: Verify socket peg data path and training smoke test
+- [ ] Run a dataset sample check against `/data/roboimi_datasets/sim_air_insert_socket_peg` with `camera_names=[l_vis,r_vis,front]` and `image_resize_shape=[224,224]`.
+- [ ] Run a short CPU or GPU training smoke test with `agent=act_resnet`, `train.max_steps=2`, `train.num_workers=0`, pretrained backbone disabled, and reduced head sizes if needed.
+- [ ] Record exact command and output snippet.
+
+### Task 5: Launch real ACT socket peg training
+- [ ] Create a run directory under `runs/` with timestamped name.
+- [ ] Start training using `/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py agent=act_resnet data.dataset_dir=/data/roboimi_datasets/sim_air_insert_socket_peg data.camera_names='[l_vis,r_vis,front]' data.image_resize_shape='[224,224]'` plus selected training hyperparameters.
+- [ ] Redirect output to `train.log` and store PID in `train.pid`.
+- [ ] Tail log to verify dataset loads, agent initializes, and first loss is produced.