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,78 @@
+# ACT Socket Peg Policy Design
+
+## Goal
+Add a local ACT-style policy/model to the existing VLA training stack and start training it on `/data/roboimi_datasets/sim_air_insert_socket_peg` using three 224×224 camera views.
+
+## Constraints
+- Base work is on branch `feat-act-socket-peg`, created from current `main` (`acbd7c605a8d203a774f2f47cff8094c05d9325e`).
+- Do not vendor or import ACT repository environment, dataset, training, or utility code.
+- Reimplement only the model/policy logic needed for this repo: CVAE action encoder, learned action queries, transformer conditioning, KL + L1 loss, and inference from prior.
+- Keep existing dataset/training loop style and checkpoint format.
+
+## Data
+The socket peg dataset is HDF5 under `/data/roboimi_datasets/sim_air_insert_socket_peg`. Episodes contain:
+- `action`: `(600, 16)`, `float32`
+- `observations/qpos`: `(600, 16)`, `float32`
+- `observations/images/l_vis`, `r_vis`, `front`: `(600, 256, 256, 3)`, `uint8`
+- attrs include `camera_names=[l_vis,r_vis,front]`, `image_height=256`, `image_width=256`, `sim=True`.
+
+Training config must use `data.camera_names='[l_vis,r_vis,front]'` and `data.image_resize_shape=[224,224]`. Existing dataset code already resizes HWC uint8 frames to `(C,224,224)` float tensors in `[0,1]`.
+
+## Architecture
+Create `roboimi/vla/agent_act.py` with `ACTAgent`, an `nn.Module` that follows the existing agent API:
+- `compute_loss(batch)` accepts `images`, `qpos`, `action`, `action_is_pad`.
+- `predict_action(images, proprioception)` returns denormalized `(B,pred_horizon,action_dim)` chunks.
+- `predict_action_chunk`, `select_action`, and queue handling mirror the existing inference contract.
+- `get_normalization_stats()` returns the current normalization module stats.
+
+Create `roboimi/vla/models/heads/act.py` containing focused, local ACT model classes:
+- Sinusoidal positional table helper.
+- Transformer encoder for posterior `z` from `[CLS, qpos, action sequence]` with padding mask.
+- Transformer decoder/action-query module conditioned on visual tokens, current qpos, and latent token.
+- `ACTPolicyHead` returning action predictions and latent `(mu, logvar)`.
+
+To avoid copying ACT's DETR image backbone code, reuse this repo's `ResNetDiffusionBackbone`. Configure it with `output_tokens_per_camera=true`, so each observation step emits one token per camera. The ACT agent builds memory tokens by concatenating camera visual tokens, a qpos token, and a latent token.
+
+## Loss
+During training:
+1. Normalize qpos/action with existing `NormalizationModule`.
+2. Keep only `num_queries == pred_horizon` actions.
+3. Encode posterior `z` from normalized current qpos and normalized target action sequence.
+4. Predict action chunk from image/qpos/latent tokens.
+5. Compute masked L1 over non-padded action timesteps.
+6. Add `kl_weight * KL(N(mu,sigma), N(0,I))`.
+
+Inference uses zero latent prior and denormalizes predicted actions.
+
+## Config
+Add `roboimi/vla/conf/agent/act_resnet.yaml`:
+- `_target_: roboimi.vla.agent_act.ACTAgent`
+- `action_dim=16`, `obs_dim=16`
+- `pred_horizon=16`, `obs_horizon=1` by default, `num_action_steps=8`
+- `camera_names: ${data.camera_names}`, `num_cams: 3`
+- ResNet backbone with `input_shape=[3,224,224]`, `output_tokens_per_camera=true`, three cameras.
+- ACT head hyperparameters small enough for a training smoke test and usable as defaults: `hidden_dim=256`, `nheads=8`, `enc_layers=4`, `dec_layers=6`, `dim_feedforward=2048`, `latent_dim=32`, `dropout=0.1`, `kl_weight=10.0`.
+
+Training command should override dataset path and camera names:
+```bash
+/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]' \
+  train.device=cuda \
+  train.num_workers=8 \
+  train.batch_size=32 \
+  train.max_steps=100000 \
+  train.use_swanlab=true \
+  train.swanlab_project=roboimi-vla \
+  train.swanlab_run_name=act-socket-peg-224-$(date +%Y%m%d-%H%M%S)
+```
+
+## Tests
+Add unit coverage without requiring ACT repo code:
+- Hydra config can instantiate `agent=act_resnet` with stubbed torchvision/diffusers-like dependencies where needed.
+- ACT loss returns a scalar, masks padded actions, and produces gradients.
+- ACT prediction returns `(B,pred_horizon,action_dim)` and honors configured camera ordering.
+- Dataset config/sample for socket peg cameras returns three resized `(obs_horizon,3,224,224)` tensors.
+