feat(vla): add ACT policy for socket peg
This commit is contained in:
Logic
@@ -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.
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user