Reference parity: `pymunk-reference` against the published Push-T

pusht_sim.envs.pymunk_reference_env.ReferencePushTEnv claims to be the environment the Push-T literature was measured on: LeRobot's gym_pusht, inherited unchanged from Diffusion Policy. This file records what that claim was checked against, what was found to differ, and what was done about each difference.

Reproduce with:

.venv/bin/python scripts/reference_parity.py

Sources read line by line on 2026-09-24:

gym_pusht/envs/pusht.py from the installed gym-pusht 0.1.6, and the same file on main. Line citations everywhere in this repo are the main numbering; the 0.1.6 file is the same code two lines further down.
diffusion_policy/env/pusht/pusht_env.py (the original the LeRobot pusht dataset was recorded in), to confirm the quirks below are not a gym-pusht regression but were there when the data was collected.

Why the harness needs two interpreters

The published environment cannot run on pymunk 7: _setup calls Space.add_collision_handler, removed in that major. This repo runs pymunk 7.3.0. So scripts/reference_parity.py runs our backend in this venv, and drives the real gym_pusht in a separate virtualenv on pymunk 6.11.1 (/tmp/refenv, --refenv to move it) through a JSON pipe.

That makes every zero below say two things at once: our reproduction matches the published environment, and the Chipmunk changes between pymunk 6 and 7 do not move this scene.

ours    python 3.12.3  pymunk 7.3.0  numpy 2.5.3
theirs  python 3.12.3  pymunk 6.11.1  numpy 2.5.3  gym-pusht 0.1.6

Measurements

Per seed the harness compares two things: the published reset(seed=n) (the published draw put through the published legacy setter), and then an open-loop rollout from one shared start with an identical action sequence generated by the harness itself. All figures are the worst deviation over every step of every episode.

Default run: 5 seeds x 40 control steps

 seed  steps        start        agent        block        angle       keypts        cover        obs32   dn  touch
-------------------------------------------------------------------------------------------------------------------
    0     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.747e-16    1.493e-05    0      8
    1     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    9.714e-17    1.525e-05    0      4
    2     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    6.939e-18    1.469e-05    0      9
    3     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.608e-16    1.510e-05    0      9
    4     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.665e-16    1.480e-05    0      6
-------------------------------------------------------------------------------------------------------------------
worst           0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.747e-16    1.525e-05    0

PASS: worst float64 deviation 3.747e-16 against tolerance 1.000e-09; worst contact-count gap 0; worst packed-observation gap 1.525e-05 against 6.200e-05 (float32)

Long horizon: 10 seeds x 120 control steps

--seeds 0,1,2,3,4,5,6,7,8,9 --steps 120. Seed 7 is in contact for 118 of its 120 steps, so this is not a measurement of two blocks sitting still.

 seed  steps        start        agent        block        angle       keypts        cover        obs32   dn  touch
-------------------------------------------------------------------------------------------------------------------
    0    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.943e-16    1.526e-05    0     14
    1    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    4.857e-16    1.526e-05    0     12
    2    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    2.602e-17    1.430e-05    0      7
    3    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.608e-16    1.523e-05    0     74
    4    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    5.551e-17    1.491e-05    0     11
    5    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.665e-16    1.509e-05    0     10
    6    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    8.674e-19    1.522e-05    0     10
    7    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.516e-05    0    118
    8    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.388e-17    1.516e-05    0     52
    9    120    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.331e-16    1.512e-05    0     18
-------------------------------------------------------------------------------------------------------------------
worst           0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    4.857e-16    1.526e-05    0

PASS: worst float64 deviation 4.857e-16 against tolerance 1.000e-09; worst contact-count gap 0; worst packed-observation gap 1.526e-05 against 6.200e-05 (float32)

The two published extra options: `--block-cog 0,60 --damping 0.9`

 seed  steps        start        agent        block        angle       keypts        cover        obs32   dn  touch
-------------------------------------------------------------------------------------------------------------------
    0     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.525e-05    0      9
    1     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.388e-16    1.525e-05    0      8
    2     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    1.469e-05    0     12
    3     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    6.072e-17    1.505e-05    0     25
    4     40    0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.608e-16    1.480e-05    0     22
-------------------------------------------------------------------------------------------------------------------
worst           0.000e+00    0.000e+00    0.000e+00    0.000e+00    0.000e+00    3.608e-16    1.525e-05    0

PASS: worst float64 deviation 3.608e-16 against tolerance 1.000e-09; worst contact-count gap 0; worst packed-observation gap 1.525e-05 against 6.200e-05 (float32)

Reading the two non-zero columns

cover, up to 4.9e-16. Coverage is the only quantity neither side reads off a body. Upstream intersects two shapely MultiPolygons built from the engine's shapes; this repo intersects the union polygon geometry.tee_polygon builds from the same vertices. Same area, different order of floating-point additions. Four ulps of a number in [0, 1].
obs32, up to 1.5e-05. PushTState.to_vector packs the observation as float32 because it crosses a JSON/HTTP boundary to the viewer; upstream's get_obs concatenates to float64. A float32 ulp at the top of the workspace is 3.05e-05, so this column is one ulp, and it is checked against its own float32 tolerance. get_state() and info["state"] keep full precision - which is what every other column is measured from - and the LeRobot pusht dataset stores float32 anyway.

Everything else - pusher position, block position, block angle, all sixteen keypoint coordinates, and the contact count - is bit-identical, at the seeded start and at every step.

Every gap found, and what was done

1. Observation angle convention - was wrong, now faithful

get_obs reports self.block.angle % (2 * np.pi) (L383); this backend reported the raw pymunk angle, so a block at -0.314 rad here read 5.969 rad there. Every LeRobot pusht state was written by get_obs, so anything comparing our state vector against the dataset saw a 2pi discontinuity.

Fixed: _poses now reduces modulo a turn, so get_state(), the packed observation and info["state"] all use the published convention. reset(state=...) still round-trips as a pose - feed -0.314 in, get 5.969 out, feed that back in and nothing moves - and the unwrapped, continuously winding angle is still on the body as env._block.angle for anything that needs it. Nothing in this repo did: every angle comparison here already goes through a wrap (experiments._wrap_angle, policies._wrap) or through cos/sin.

2. Initial-state sampling - was wrong, now faithful

Upstream draws integers(50, 450), integers(50, 450), integers(100, 400), integers(100, 400), uniform(-pi, pi) (L276-280) from the generator gymnasium builds as np.random.default_rng(seed). BasePushTEnv.sample_state drew uniform floats over 50..462 and 100..412 - a different support, so a seeded start here was not a start the published environment could produce.

Fixed: ReferencePushTEnv.sample_state overrides the base draw with the published one. Same constructor, same order, so seed n here selects the state seed n selects there - the start column above is that equality, measured end to end at 0.0.

Consequence, deliberate: this is the one place the backend breaks the base class's "one seed, one scenario on every backend" rule, so make_env("pymunk", seed=n) and make_env("pymunk-reference", seed=n) no longer agree on where the episode starts. Anything comparing engines has to sample once and pass the state to the others with reset(state=...); the published draw is integral and strictly inside the SI backends' support, so it is always a legal start for them. The web session does exactly this.

3. Legacy `_set_state` - was a deviation, now a deviation with an exact entry point

Upstream assigns the block position and then the angle (L457-462) and says why: "for compatibility with legacy data". pymunk rotates about the centre of gravity, which for this block is (0, 45), so the pose asked for is not the pose obtained - upstream's own docstring shows [0, 10, 20, 30, 1] coming back as (57.87, 50.69), a shift of 42 units. Recorded LeRobot/Diffusion Policy episode starts were all written through that setter, so replaying one through a setter that honours the pose puts the block in the wrong place.

Both behaviours, explicitly:

reset(state=...) keeps add_tee's angle-first order (L510-511), so the pose round-trips. This is what PushTState, geometry.transform_points and both other backends mean by a pose, and the whole three-way comparison rests on it.
set_reference_state(vector_or_state) is the legacy setter, verbatim, and reproduces the documented shift exactly. Use it for recorded data.
reset_reference(seed) is the two published steps together - published draw, legacy setter - i.e. gym-pusht's reset(seed=n) in one call.

4. Missing accessors - added

get_keypoints() (L517-537): the eight block corners in world coordinates, (8, 2); .flatten() is the 16-value environment_state LeRobot trains environment_state_agent_pos policies on. Taken from the engine's own shapes, because the index order in upstream's diagram is Chipmunk's re-winding of the input vertices and not geometry.tee_vertices' order. Both pymunk majors were checked to return the same winding.
Contact count (L425-426, L413-420): info["n_contacts"], on pymunk 7's Space.on_collision(0, 0, post_solve=...) in place of the removed add_collision_handler. Only post_solve is registered, so the solver is untouched - and the parity runs above, which include an episode in contact for 118 of 120 steps, are the evidence that registering it changed no physics.
block_cog and damping: constructor arguments of ReferencePushTEnv, exactly as they are constructor arguments of PushTEnv (L140-141, L165-166), and measured to be exact above. They are deliberately not in SimParams: that dataclass is the shared, swept, serialised schema of the two SI engines, and neither of them has a "move the block's centre of gravity" or "multiply every velocity by this each step" knob to map them onto. Putting them there would add two permanently-inapplicable entries to every schema response, every sweep and every UI form. make_env therefore keeps offering the published defaults, and a caller who wants the options constructs the class.

5. Found while reading the two sources side by side

info["n_contacts"] is divided by the wrong number (L414). _get_info computes n_steps = int(1 / self.dt * self.control_hz), which is 1000, where step (L240) computes the substep count as int(1 / (self.dt * self.control_hz)), which is 10. So the published "contact points per step" is the real count over a hundred times too large a divisor: it is 1 whenever anything touched during the step and 0 otherwise. Reproduced bug-for-bug, with the raw sum published alongside it as info["n_contact_points"].
get_keypoints is broken on NumPy 2. L537 calls np.row_stack, removed in NumPy 2.0, so obs_type="environment_state_agent_pos" raises AttributeError on the pinned gym-pusht 0.1.6 + NumPy 2.5.3 install. Deviation: ours uses np.vstack, which is what row_stack was an alias for - same values, and it runs. (The parity harness restores the alias in the reference subprocess so the published function can be called at all.)
No action clipping upstream. step declares Box(0, 512) but never enforces it; BasePushTEnv.step clips. Inside the declared range - every action a policy or a dataset produces - the two are identical. Left as is.
Observation dtype. Upstream concatenates to float64; PushTState.to_vector packs float32. See "Reading the two non-zero columns" above.
Termination comparison. Upstream terminates on coverage > 0.95 (L256, strict); the base class uses >=. They differ only on an episode sitting exactly on the threshold.
Truncation. Upstream never truncates (L263 truncated = False); the base truncates at max_episode_steps. Base-owned, documented, left as is.
The two physics quirks were already reproduced and still are: shape friction is 0.0 everywhere because body.friction = 1 (L476, L512) assigns to a property pymunk.Body does not have, and the block's moment is twice the bar's because L501 passes vertices1 where it means vertices2. Both are present in the Diffusion Policy original too, so the recorded dataset has them, and tests/test_env_reference.py pins both.

Tests

tests/test_env_reference.py covers all of the above without needing the reference virtualenv - gym-pusht's semantics are re-stated there as literals from its source rather than imported. The single test that runs the real package shells out to scripts/reference_parity.py and skips when /tmp/refenv is absent.

.venv/bin/python -m pytest tests/test_env_reference.py -q

Reference parity: pymunk-reference against the published Push-T