LeWM parity: lewm against LeWorldModel's Push-T

pusht_sim.envs.lewm_env.LewmPushTEnv claims to be the environment LeWorldModel trains and evaluates on: stable_worldmodel.envs.pusht.PushT, registered as swm/PushT-v1 in galilai-group/stable-worldmodel, whose built-in WeakPolicy generated the lewm-pusht / pusht_expert_train dataset. This file records what that claim was checked against, what was found to differ from the repo's existing contracts, and what was done about each difference. It is the counterpart of reference_parity.md, which does the same for gym-pusht: the two environments share a pymunk scene - swm/PushT-v1 is a fork - and nothing else.

Reproduce with:

.venv/bin/python scripts/lewm_parity.py --seeds 0,1,2,3,4 --steps 200

Sources read line by line on 2026-09-24, at commit 63988116d34cde56aea1240d5e58eb158ac67dc0 (installed as stable-worldmodel 0.1.1 in /tmp/swmenv):

Why the harness needs two interpreters

Not because of pymunk this time: both sides run pymunk 7.3.0 and NumPy 2.5.3. stable-worldmodel pulls in torch, pygame, an environment registry and a training stack, none of which belongs in this repo's venv, so the published environment is driven in its own virtualenv (/tmp/swmenv, --swmenv to move it) through a JSON pipe by scripts/lewm_parity.py.

ours    python 3.12.3  pymunk 7.3.0  numpy 2.5.3
theirs  python 3.12.3  pymunk 7.3.0  numpy 2.5.3  stable-worldmodel 0.1.1

Because the engine is the same on both sides, a zero below says one thing only: this class, its seed expansion, its action translation and its policy are the published ones.

What is compared

Per seed, one closed-loop episode. reset(seed=n) on both sides, then each side runs its own seeded WeakPolicy(dist_constraint=100, seed=n) - upstream's from the package through its vectorised float32 path (the one that generated the dataset; the single-env path is broken upstream), ours from pusht_sim.policies through LewmPushTEnv.command - each reading the state of its own environment. A physics disagreement would therefore compound through the policy, so a zero means the episode is identical action for action. Every step's 7-D observation, reward, success flag and contact count is compared, plus the reset's observation and info["goal_state"]. All figures are the worst deviation over every step of every episode.

Default run: 5 seeds x 200 control steps

 seed  steps       start        goal       agent       block       angle         vel      reward  dsucc   dn  succ  touch
-------------------------------------------------------------------------------------------------------------------------
    0    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0     85
    1    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0     77
    2    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    103
    3    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    146
    4    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    140
-------------------------------------------------------------------------------------------------------------------------
worst          0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0

  start   the 7-D observation after reset(seed=n)     goal    info['goal_state'], all 7 components
  agent   pusher position, px, over every step       block   block position, px
  angle   block angle, radians (both sides wrapped)   vel     pusher velocity, px/s
  reward  -|goal_state - state| over 7 components     dsucc   steps whose success flag differs
  dn      worst gap in info['n_contacts']             succ    steps upstream reported success
  touch   steps of this episode in which upstream reported a contact

PASS: worst float64 deviation 0.000e+00 against tolerance 1.000e-09; worst contact-count gap 0; success flag mismatches 0

Contact stress: 20 seeds x 200 control steps, --dist-constraint 30

A 30 px box keeps the pusher on the block: seeds 13 and 17 are in contact for all 200 steps, and no episode is out of contact for more than a third of its length. This is the run that exercises the solver rather than the PD law.

 seed  steps       start        goal       agent       block       angle         vel      reward  dsucc   dn  succ  touch
-------------------------------------------------------------------------------------------------------------------------
    0    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    195
    1    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    178
    2    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    199
    3    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    182
    4    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    193
    5    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    183
    6    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    170
    7    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    154
    8    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    196
    9    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    174
   10    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    173
   11    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    159
   12    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    174
   13    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    200
   14    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    182
   15    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    134
   16    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    163
   17    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    200
   18    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    181
   19    200   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0     0    193
-------------------------------------------------------------------------------------------------------------------------
worst          0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00   0.000e+00      0    0

PASS: worst float64 deviation 0.000e+00 against tolerance 1.000e-09; worst contact-count gap 0; success flag mismatches 0

Reading the table

Every column is exactly zero: the start observation, the sampled goal state, all seven components of every step's observation, the reward, the success flag and the contact count agree to the bit over 200 steps of closed loop. The succ column is 0 on every row because WeakPolicy never reaches the goal in 200 steps - it is a random walk around the block - so the success flag's true branch is pinned by tests/test_env_lewm.py against a threshold state measured on the installed package instead.

The harness is sensitive. Feeding the same rollout float64 actions instead of the published float32 ones (see gap 1 below) moves the block by up to 8.8e-05 px and the reward by 1.7e-04 over 200 steps of seed 0 - five orders of magnitude above the 1e-09 tolerance.

Every gap found, and what was done

1. The action's dtype is part of the physics - reproduced

The published action space is Box(-1, 1, (2,), float32) and the dataset generator fills a np.float32 buffer (expert_policy.py L63). In step, action = self.agent.position + action * self.action_scale (L325) adds a float32 array to a Vec2d of Python floats, and under NumPy 2's promotion rules the result is a Vec2d of np.float32. Every operation of the PD law - k_p * (action - position), + k_v * (0 - velocity), * dt, velocity += - then rounds to single precision before pymunk stores the velocity as a double. Measured on the installed package: after one control step from reset(seed=0) the pusher velocity is exactly float32-representable with a float32 action and ~6e-05 px/s different with a float64 one.

Reproduced: LewmPushTEnv.step keeps the dtype it is handed and spells the arithmetic with the same operands in the same order (Vec2d + ndarray, int * Vec2d, ...). command emits float32, because that is the published action dtype and what the dataset generator produced. A float64 action to step runs the same law in double precision. Pinned by test_action_dtype_is_part_of_the_arithmetic.

2. The seed expands through gymnasium's Dict.seed, with sorted keys - reproduced

Upstream never draws from a single generator. reset seeds the nested variation Dict with the reset seed (spaces.py L40); gymnasium's Dict.seed(int) draws integers(2**31 - 1, size=len(children)) from default_rng(seed) and hands one subseed to each child in key order - and gymnasium sorts a Dict's keys at construction, so the order is agent, background, block, goal, rendering at the top, not the order env.py lists them. Nested dicts repeat the process; each leaf Box becomes its own default_rng(subseed). The three default variations are sampled once into .value (the start) and once more with set_value=False (the goal, L268-283). The first attempt at this, in env.py's listing order, put the block in the wrong place; the sorted order reproduces the published numbers exactly.

Reproduced: lewm_env.sample_variations(seed), numpy only, and reset(seed=n) feeds n straight into it. reset() without a seed draws a stream seed from the env's own generator where upstream would use OS entropy, so an unseeded reset stays reproducible from the constructor seed. Pinned against reset(seed=0)'s start and goal to full precision, and against a hand re-derivation of the subseed tree.

3. The goal pose leaks into the start pose - reproduced

reset poses the scene at the goal, steps once, renders the goal image, then poses it at the start and steps once more (L286-307) - in one space. block.start_position runs to 400 while the T reaches 120 px from its origin, so a sampled goal block is often through a wall. Chipmunk pushes it out partly through a bias velocity it keeps on the body and applies in the next step's position update; _set_state overwrites position and angle but not that bias, so the start's settle step moves the block by whatever the goal pose's overlap left behind, with no contact at the start pose at all. block.velocity reads zero throughout (space.damping = 0 wipes it), so this is invisible from the state. Seed 8 moves 0.009 px, seed 23 0.28 px, seeds 17 and 32 (whose start poses touch too) 0.71 and 0.55 px, before the first action.

Reproduced by running the same two _set_state calls in one space, which is also why reset(state=...) still performs the goal settle. The reset's n_contacts counts both settle steps (nothing clears the counter between them). Pinned on seed 17.

4. Relative actions on a repo built around absolute targets - one hook

Every policy in pusht_sim.policies emits an absolute target and every driver - rollout, the web session, the drag - called env.step(target). swm/PushT-v1 takes a displacement in [-1, 1] scaled by 100 px. Making step accept both would mean guessing from the values.

Done: BasePushTEnv.command(target), which is step(target) on every engine whose action is the target, and on lewm is step(clip((target - agent) / 100, -1, 1).astype(float32)) - the exact rescaling WeakPolicy performs upstream before handing its action to the env (expert_policy.py L79-80). rollout, Session.step (which the drag and the policy run both go through) call command; step keeps each engine's published contract. latest_action on lewm is the absolute target the action resolved to, so the viewer's crosshair still shows where the pusher is being told to go.

WeakPolicy in this repo returns the clipped point as the target (upstream's value one line before its rescaling), so command(act(state)) hands the env the bit-identical float32 action the dataset generator did - the table above is that equality - and the same target drives the other three engines through their own command. default_policy_kwargs("lewm") gives the scripted policy the position-mode knobs the kinematic pusher wants, the same as pymunk-reference.

5. Two goals: the scored state and the drawn overlay - both exposed, named apart

Upstream has goal_state, the sampled 7-vector eval_state scores against, and goal_pose, the goal.position / goal.angle variation (default (256, 256, pi/4)) that is only ever drawn. The repo's PushTConfig.goal_pose and PushTState.goal_* were one thing.

Done: on lewm, PushTState.goal_* (and so the renderer, the web screen's geometry and coverage()) is the sampled goal block pose - the thing that is scored - and info["goal_pose"] is config.goal_pose, the overlay, exactly as upstream names them. env.goal_state / info["goal_state"] carry the full 7-vector. reset(state=s) takes s.goal_* as the goal block pose and the seed's own draw for the goal pusher position (there is no other source for it); reset(goal_state=...) overrides all seven, upstream's options["goal_state"]. Coverage remains available as a secondary metric, against the sampled goal; reward, terminated and info["is_success"] follow eval_state.

6. Found while reading the two sources side by side

Tests

tests/test_env_lewm.py covers all of the above without needing the stable-worldmodel virtualenv - the environment's semantics are re-stated there as literals from its source and its numbers as values measured on the installed package. The single test that runs the real package shells out to scripts/lewm_parity.py and skips when /tmp/swmenv is absent.

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