IXDLI/pusht-parallel-rain-seed0

PushT rain, seed 0

Final EMA policy trained from scratch for 512,400 optimizer updates. RGB96 + agent XY, two observations, 16 prediction / eight execution steps. Original lossless DP PushT archive; original 90-demo split; effective batch64 (final epoch batch may be smaller), AdamW1e-4, cosine500 warmup, warmed EMA. The final checkpoint is published regardless of test score; it is not test-selected.

Method and reproduction scope

• RGB PushT controlled reproduction; original DP 90-demo split and normalization, fixed 512400 optimizer updates at effective batch64.
• TDP is derived from the paper, as authors have not released code or complete PushT hyperparameters. Tube constants and action-to-U-Net lift are explicit implementation choices.
• GRN is an action-domain adaptation, with a larger configurable temporal head.
• DiT uses Dasari et al. DiT-Block conditioning adapted to shared PushT RGB encoder, not a pretrained VLA.
• Modern gym-pusht preserves legacy state reset order but is not the original archived DP environment; version recorded.
• Periodic checkpoint evaluation reports test score as in the paper; final published policy is the final EMA, not the best test-selected checkpoint.
• RAIN RGB2D PushT adaptation of official uncertain ACT + action-to-action FM DiT + calibrated risk ensemble; original paper ManiSkill state PushT is a different environment.
• User-selected RAIN-DP changes only the refiner backbone to the DP conditional U-Net512/1024/2048; retains the RAIN flow-matching objective, not a DDPM/DDIM loss.
• Total512400 optimizer updates are shared across proposal126000, frozen-proposal refiner126000, joint252000, and frozen-backbone risk ensemble8400. Refiner and actor each receive378000 updates, not512400 each.
• Shared original DP90 training demonstrations and min/max normalization; the original four initial validation episodes alone calibrate the router. Other112 train-complement demonstrations are unused.
• Common batch64, AdamW1e-4 betas.95/.999 weightdecay1e-6 eps1e-8. RAIN uses stage-local cosine/warmup, differing from the single-stage DP schedule.
• ACT prediction queries16 rather than original30, common Ho2/Hp16/Ha8. Both XY dimensions are continuous and refinable; there is no gripper dimension.
• Causal temporal aggregation uses the previous H8 query for only the current action; cache targets and calibration use the exact deployed source and H8 stride, correcting released C0 and H2/H8 mismatches.
• Offline frozen-EMA candidates0/1/2/3 supervise the three-member risk ensemble. No reward/test metrics tune routing. Before router calibration periodic evaluations use fixed3 refinements and label this explicitly.
• Final EMA proposal/refiner plus the risk ensemble selected on held-out demonstration loss and offline calibrated coefficients are published; no test-score model selection.

Method-specific settings and any controlled-comparison adaptations are explicit in config.json. For RAIN, the total budget includes its four training stages; proposal/refiner backbone updates and auxiliary router updates are reported separately. This controlled RGB comparison does not assert reproduction of reported paper scores. Evaluation JSON reports mean maximum reward separately from strict success rate.

Load

Install requirements.txt using a CUDA-compatible PyTorch/torchvision pair, extract inference_source.tar.gz, and run:

from pusht_parallel.inference import load_policy, PushTController
model, payload = load_policy("policy.pt", device="cuda")
controller = PushTController(model, payload, device="cuda", nfe=10)
action_xy = controller(observation)  # {"pixels": uint8[96,96,3], "agent_pos": float[2]}

Call controller.reset() before each episode. Model outputs are XY position commands.

Sources: Tube DP, original DP experiment, DiT-Block Policy, GRN, RAIN.

W&B: https://wandb.ai/yuyuan-wu-universiteit-gent/policy%20parallel%20test/runs/pusht-parallel-20261019-s0-rain-s0