SleepMastger/srlwam-cosmos-edge-stack-v2-stride4

srlwam: Cosmos Edge stack_v2 sim policy, stride 4

A Cosmos3 Edge policy for Franka two-cube stacking, fine-tuned from nvidia/Cosmos3-Edge-Policy-DROID on the stack_v2 simulation demos with a temporal stride of 4: each model step covers 4 recorded 10 Hz steps. It is the stride-4 variant of the sim_v2_droid checkpoints in aabyaneh/srlwam, built from Amin's srlwam cosmos_stack_v2_stride3 code with the stride changed to 4.

checkpoint	step	held-out normalized action MAE	gripper accuracy	flow loss
sim_v2_stride4_droid_iter1750	1750 (best)	0.0525	0.988	0.72
sim_v2_stride4_droid_iter2000	2000 (final)	0.0537	0.987	0.71

Validation used 744 fixed windows from the 250 held-out sim episodes, 30 sampling steps, open-loop against the recorded actions. These are not robot or closed-loop success rates. For reference, the stride-1 sim_v2_droid run (same data and recipe) reached 0.0593 at step 1750 and 0.0624 at step 2000. Each run normalizes its own action space, so this is only a rough comparison.

What stride 4 changes

• Timing: one model step is 4 recorded steps (0.4 s, 2.5 Hz). A 16-step action chunk spans 6.4 s of motion.
• Actions: each action is the sum of 4 recorded deltas (translation and rotation), with the gripper command of the last of them. The normalization was refit on these summed actions: xyz q01/q99 is about ±0.04–0.06 sim units, against ±0.01–0.016 at stride 1.
• Video and state: sampled every 4th recorded frame. conditioning_fps is 2 (2.5 Hz floored, because the model only accepts an integer).
• Everything else matches the stride-1 sim run: 750 train / 250 held-out demos, 2000 steps × global batch 2048, lr 5e-5 with 500 warm-up steps, 8×H200, bf16, init from Cosmos3-Edge-Policy-DROID.

Deploying: two changes from the stack_v1/v2 deploy code

This repo does not include deploy code, because the shared srlwam Cosmos adapter (deploy/cosmos/cosmos_policy.py in aabyaneh/srlwam) assumes stride 1. To run these weights with it:

1. Set FPS = 2 (not 10) in cosmos_policy.py, so conditioning_fps matches training.
2. Execute each predicted action over 0.4 s: either as one 2.5 Hz step, or as 4 equal 10 Hz sub-steps of action / 4. Do not play it as a single 10 Hz step, which would move 4× too fast.

Everything else follows deploy.json. The observation frame is sim_ee_frame, because the policy was trained on sim state (convert raw robot readings as the adapter does). Units are stack_v2 sim units at gain 0.1975 m/unit, so use command.delta, not real_controller_delta. The adapter's 0.025 m per-step clamp sits above typical summed steps (q99 ≈ 0.013 m). Sim-only stack_v2 checkpoints have shown real-state range and descent problems on the real robot data (see aabyaneh/srlwam), so start in shadow mode.

No fixture_reference.npz or eval_real_holdout.json is included, so these checkpoints have not been through gate-1 or real-data evaluation. They have not been evaluated on a robot.

Layout

cosmos_edge/stack_v2_stride4/<checkpoint>/
├── model/                  bf16 DCP (.metadata + __0_0.distcp) + config.json (non-EMA net.* weights)
├── deploy.json             obs_frame, stride, action units/gain, training state range, sha256, lineage
├── checkpoint.json         policy metadata (16-step chunk, conditioning_fps 2, stride 4, domain 33, "stack the cubes")
├── normalization.json      stride-4 action normalization (refit on summed actions) + state normalization
├── split.json              sim train / held-out episode split
└── train_validation/       metrics at this step, validation error plot, resolved training config
cosmos_edge/shared/         Wan2.2 VAE, processor/tokenizer files, vision tower (identical to aabyaneh/srlwam)
training_code/              the exact stride-4 training tree (stack_policy/, scripts/, configs/, readiness receipt, Slurm wrapper)
tools/                      export.py (training DCP -> these bundles), upload.py
SHA256SUMS.txt              sha256 of every file in this repo

Training: Slurm job 5050 (wandb huaweiwam/srlwam 6ha0xwzr). Data: dataset repo aabyaneh/srlwam @ cffcf593 (stack/sim/stack_v2).

Licenses

These are fine-tuned derivatives of nvidia/Cosmos3-Edge-Policy-DROID (OpenMDW 1.1). The shared Wan2.2 VAE is Apache 2.0. See cosmos_edge/NOTICE.md.