IXDLI/pusht-parallel-grn-seed0

PushT grn, 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.

TDP settings not supplied by its authors are explicit in config.json, including contraction/noise constants, continuous interpolation, and single-action U-Net lift. 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.

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