File size: 2,530 Bytes
2d7bd36
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
---
tags:
- robotics
- imitation-learning
- pusht
- diffusion-policy
---
# PushT dit, 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:

```python
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](https://arxiv.org/abs/2604.23609),
[original DP experiment](https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/config.yaml),
[DiT-Block Policy](https://github.com/SudeepDasari/dit-policy),
[GRN](https://github.com/bytedance/GRN).

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