AaronHuangWei's picture
Add files using upload-large-folder tool
aab194c verified
|
Raw History Blame Contribute Delete
5.44 kB
---
license: other
pipeline_tag: robotics
tags: [robotics, world-action-model, fast-wam, baseline, unitree-g1, real-robot]
extra_gated_prompt: "Private real-robot baseline bundle (Fast-WAM on the Unitree G1, ablation for Long-WAM). Access is granted manually by the owner."
extra_gated_fields:
Name: text
Affiliation: text
Intended use: text
---
# Fast-WAM baseline on the Unitree G1 (ablation vs Long-WAM)
Same two-process layout and the **same robot-side code** as the Long-WAM deployment
(`AaronHuangWei/Long-WAM-G1-Dynamic-Task-Deploy`): the policy runs on the workstation GPU behind a
WebSocket, the robot PC streams observations and applies joint targets. Only the server-side model
differs.
```
robot PC (eval_longwam_g1.py) --30 Hz obs (3 JPEG + 16-D state)--> workstation (policy_server.py --arch fastwam)
<-- 32 x 16 joint targets every replan -- FastWAMPolicy (plain BF16, no history)
```
## What this baseline is
| | Long-WAM (our method) | Fast-WAM baseline (this bundle) |
|---|---|---|
| model | ARWAM P4: LongLive-2.0 AR video expert + ActionDiT | original Fast-WAM: Wan2.2-TI2V-5B DiT + interpolated ActionDiT |
| init / pretraining | YAM 130k-episode robot pretraining, fine-tuned per task | Wan2.2 base weights, trained directly on each task's demos (Fast-WAM default recipe) |
| observation at inference | current frame + 12 history frames (1.6 s, stride 4) | **current frame only** (no history, no memory) |
| future imagination | 2 future latents denoised to sigma 0.9, actions read them | **none**: `infer_action` denoises the action chunk directly from the current-frame latent |
| action denoising | 10 flow-matching steps | 10 flow-matching steps (Fast-WAM's own eval default) |
| inference acceleration | resident RoPE, torch.compile, CUDA Graphs (lossless) | **none** – plain BF16 eager PyTorch, as released |
| training | 20K steps, global batch 512 | 5K steps, global batch 512 (same data, cameras, dims, normalization, prompts) |
Cameras, action dims, normalization statistics and prompts are identical to the Long-WAM runs
(`contract.TASK_SPECS`): Dynamic_Cup = 3 cameras / 16-D, Speed_Cup_20/30/40/50 = 2 cameras / 8-D.
Actions are absolute joint targets, 32 steps per chunk at 30 Hz, returned in the full 16-D robot
order with non-controlled dims held at the current state.
## Files
| file | runs on | purpose |
|---|---|---|
| `policy_server.py --arch fastwam` | workstation | loads the checkpoint, serves WebSocket (`observe` requests are accepted and ignored: no history) |
| `fastwam_policy.py` | workstation | Fast-WAM wrapper: pinzi frame (identical to training), normalization, `infer_action` |
| `longwam_policy.py`, `contract.py` | workstation | shared frame composition and the fixed I/O contract |
| `eval_longwam_g1.py`, `msgpack_numpy.py`, `websocket_client.py`, `contract.py` (`robot/`) | robot PC | unchanged from the Long-WAM bundle |
## Bundle layout
```
README.md
robot/ the 4 files the robot PC needs (no torch)
code/FastWAM/{src,deploy/unitree_g1,pyproject.toml}
models/<task>/config.yaml training config of the Fast-WAM run (model=fastwam)
models/<task>/step_005000.pt final weights (5K steps)
models/<task>/fastwam_dataset_stats.json normalization (same file as the Long-WAM run of the task)
models/Wan-AI/Wan2.2-TI2V-5B/{Wan2.2_VAE.pth,config.json}
text_embeds_cache/ cached T5 embeddings of the two prompts (avoids the 11 GB T5)
```
## Workstation
Requirements: Python 3.10/3.11, PyTorch 2.7.1+cu128, then `pip install -e code/FastWAM websockets msgpack`.
```bash
cd bundle/code/FastWAM
export DIFFSYNTH_MODEL_BASE_PATH=$PWD/../../models # contains Wan-AI/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
export PYTHONPATH=src:deploy/unitree_g1
python deploy/unitree_g1/policy_server.py --arch fastwam \
--run-dir ../../models/Speed_Cup_20 \
--checkpoint ../../models/Speed_Cup_20/step_005000.pt \
--stats ../../models/Speed_Cup_20/fastwam_dataset_stats.json \
--text-cache-dir ../../text_embeds_cache \
--task Speed_Cup_20 --port 8000
```
`--accel`, `--compile-mode` and `--video-steps` are ignored for `--arch fastwam` (Fast-WAM runs
eager on purpose; the Long-WAM accelerations are not applied to the baseline). `--action-steps`
defaults to 10. Start-up is about one minute (no compilation).
## Robot PC
Identical to the Long-WAM bundle:
```bash
cd unitree_lerobot
PYTHONPATH=/path/to/bundle/robot:$PYTHONPATH python /path/to/bundle/robot/eval_longwam_g1.py \
--policy_server_host WORKSTATION_IP --policy_server_port 8000 \
--arm G1_29 --ee dex1 --control_side right --replan_steps 8 \
--send_real_robot true --motion true
```
Use `--control_side both` for `Dynamic_Cup`, `right` for the `Speed_Cup_*` checkpoints. Keep the
same `--replan_steps` as in the Long-WAM runs so the two methods are compared under the same
control schedule.
## What was tested
- H100 loopback (`loopback_g1_policy.sbatch` with `ARCH=fastwam`): `policy_server.py --arch fastwam`
driven by the synthetic JPEG client through the real WebSocket path for Speed_Cup_20 (2 cameras,
8-D) and Dynamic_Cup (3 cameras, 16-D); output shape 32 x 16, finite, non-controlled dims held.
Latency numbers are in the training-side log; they are not part of the comparison (the baseline
is intentionally unaccelerated).
- Not tested here: the real Unitree hardware loop.