roboimi/docs/superpowers/plans/2026-05-08-act-socket-peg.md

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.