refactor(sim): give agents their own interface so the env stops branching on them

[ElmoPA]

Adding the u_socket was done by editing the environment. Its latch, friction and
penetration guards -- 12 methods, ~470 lines -- went into env.py as
if pusher_shape == "u_socket" branches, and its 3-DOF action became a hardcoded
expected_shape = (3,) if self.pusher_shape == "u_socket" else (2,). That is the
single largest reason the two sims diverged: sim_v1's env.py has ZERO socket
references, sim_v2's had 115, and env.py's step loop called nine socket-specific
guards in sequence.

An Agent now owns the three things the environment should not know about:

• ACTION SPACE -- action_dim (2 for a free-moving pusher, 3 when the agent
  also controls orientation) and target_pose() to decode a raw action;
• BODY -- build() in the pymunk space;
• CONTACT MODEL -- pre_substep()/post_substep() hooks around each physics
  substep, plus on_reset() for per-episode state.

env.step is agent-agnostic:

captured = self.agent.pre_substep(self)
self._drive_pusher_toward(tx, ty, dt_sub, target_angle)
self._space.step(dt_sub)
self._clamp_pusher_to_static()
self.agent.post_substep(self, captured)

Agent (circle, circle_small, stick, L) implements the hooks as no-ops.
USocketAgent owns all the latch/guard logic and the socket geometry constants,
and its solid_pusher / socket_inside_friction_only flags become constructor
arguments rather than environment state. A new agent with an unusual action
space is a new class plus one line in make_agent(), not another branch in the
simulator.

env.py 1292 -> 795 lines.

sim_v1 IS DELIBERATELY UNTOUCHED. It is frozen so pre-rewrite data replays
exactly; refactoring it would put that at risk for no benefit, since it has no
socket to abstract in the first place.

VERIFIED BY REPLAY EQUIVALENCE, not by inspection. Baselined the unmodified sim
with the identical harness first, then compared:

u_socket_3000_v2         100.0% -> 100.0%   (p50 0.0033 -> 0.0033)
circle_3000_plus_gen_v2   89.7% ->  89.7%
circle_v2_obstonly        17.1% ->  25.7%

The gate caught two real bugs that inspection did not: the moved
_socket_contact_is_on_inner_face call site lost its env argument, and
socket_latched -- a property DERIVED from _socket_constraints is not None --
had become a plain attribute nothing updated, initialised to [] so it would have
read as permanently latched. Both fixed; the socket went 0% -> 100%.

Co-Authored-By: Claude Opus 5 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012V58H37tmcvgDthELMd5Xk

[ElmoPA]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• feat(configs): hydra config groups for the batchflow stack #558 : 2 dependent PRs (#560 , #564 )
• test: batchflow test suite and retained tooling #557
• feat(eval): batchflow evaluation stack #556
• feat(algo): H-Net, diffusion and BC algorithms #555
• feat(models): H-Net and diffusion model zoo #554
• refactor(sim): give agents their own interface so the env stops branching on them #559 👈 (View in Graphite)
• sim: Tsimulation with sim_v1 and sim_v2 #553
• feat(infra): batchflow pipeline core, rldb unions, collapsed embodiment enum #552
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Claude Code Review

Review

Summary

Extracts pusher-specific logic (action space, body build, contact model) from env.py into an Agent hierarchy, with USocketAgent owning the ~470 lines of socket latch/friction/guard code. Verified by replay equivalence against a baseline of the untouched sim.

Key concerns

1. PR scope creep. The title and description are strictly about the agent refactor, but the diff also adds --speed-factor, --pusher-color, --variant-name to mouse_collect.py and a new metadata_override parameter to ZarrDemoWriter. These are unrelated features and should be in a separate PR — mixing them dilutes the replay-equivalence gate that gives this refactor its credibility, and makes it harder to bisect later.

2. env.PUSHER_SPEED = type(env).PUSHER_SPEED * float(args.speed_factor) — this pattern in mouse_collect.py mutates instance attributes shadowing class attributes right after construction. If any internal code reads PushShapesEnv.PUSHER_SPEED (class-level) rather than self.PUSHER_SPEED, speeds won't scale. This also isn't recorded in episode_init, so replay of these episodes will not reproduce the actual dynamics — only the metadata_override records it. That's a real replay divergence hazard, exactly the class of bug the refactor is meant to prevent.

3. Global mutation of render_module.PUSHER_COLOR. Setting a module-level global based on CLI args means any test or concurrent process importing the same module gets a stale/mutated color. Should be per-env state, not a module global.

4. Lambda in on_collision handler:

   pre_solve=lambda a, sp, d: self.agent._socket_friction_pre_solve(self, a, sp, d)

   This closes over self, keeping the env alive as long as the space handler is registered. Since _space is rebuilt per episode this is probably OK in practice, but the two smoke tests had to be modified to inject a monkey-patched handler with a different signature (agent_env, arbiter, space, data) — which means the test is no longer exercising the real registered callback. Consider a bound method or partial and keep the test signature aligned with the real one.

5. _socket_friction_pre_solve is prefixed private but called from env via self.agent._socket_friction_pre_solve(...). If the intent is that the env drives contact-model hooks generically, this should probably be a public on_collision_pre_solve on the base Agent (no-op by default) rather than an env branch that reaches into an agent-private method. Otherwise you've re-introduced a u_socket-shaped coupling one layer up (if getattr(self.agent, "socket_inside_friction_only", False)).

6. get_episode_init no longer includes agent-specific fields. Agent.episode_init() was added with an init_fields mechanism but I don't see env.get_episode_init() updated to call it. If solid_pusher / socket_inside_friction_only used to be captured for replay via env attrs, and now live on the agent, verify the episode_init dict still contains them — otherwise older replay code that reads these fields could silently get defaults.

Suggestions

• Split the color/speed/variant CLI additions and ZarrDemoWriter.metadata_override into a follow-up PR. Keep this one purely the refactor.
• If you keep them here: fold speed_factor into env_args / episode_init properly (e.g. an env constructor kwarg) rather than post-construction attribute mutation, so replay reproduces the physics.
• Replace the render_module.PUSHER_COLOR = ... global assignment with a per-env override.
• Promote _socket_friction_pre_solve and the socket_inside_friction_only check to a generic Agent.register_collision_handlers(env) hook so env.reset has zero agent-type awareness. Right now the env still says "if the agent has this socket-specific flag, register this socket-specific handler."
• Add a test that env.get_episode_init() round-trips solid_pusher and socket_inside_friction_only post-refactor.
• Consider a test that instantiating an unknown pusher raises the ValueError from make_agent.
• The Agent base class does self.solid_pusher = bool(solid_pusher) unconditionally, but init_fields = () for the base class — so simple agents don't surface solid_pusher in episode_init even though they respect it. Either put solid_pusher in the base init_fields or document why not.

Verdict: Request Changes

The core refactor is well-motivated, well-executed, and the replay-equivalence gate + the two bugs it caught are exactly the right way to land a change like this. But the PR bundles an unrelated feature (speed/color/variant) with a physics-affecting attribute mutation that is not captured in episode_init — precisely the kind of hidden divergence this refactor exists to prevent. Please split.

---

Reviewed by Claude · Review workflow