Hook system for simulation callbacks

[juan-g-bonilla]

Summary

Replaces the hard-coded step-loop side effects scattered across the simulator backends with a single lifecycle hook registry. Instead of each backend hand-calling virtual_gantry.step(), _step_bridge(), and capture_video_frame() inside its simulate_at_each_physics_step, and overriding on_episode_start/end, the runtime participants now register callbacks against typed Phase events that BaseTask (and the run_sim DirectSimulation loop) emit.

1. hooks.py — new HookRegistry + Phase enum. Seven lifecycle phases: FRAME_BEGIN / FRAME_END bracket the outer control-rate tick, PRE_STEP / POST_STEP bracket each physics substep, plus EPISODE_START, EPISODE_END, CLOSE. Deterministic registration-order dispatch, per-phase payload-arity validation, enable/disable/remove handles, mutation-during-emit guards, and a reverse-order CLOSE that runs every hook exactly once and aggregates failures into one HookCloseError. Covered by tests/test_hooks.py (no_sim).
2. Per-hook cadence — every. hooks.add(phase, fn, every=4) runs a callback once per N emissions of its phase. On periodic phases (FRAME_*, *_STEP), every also accepts a frequency string ("30Hz", ">30Hz", "<30Hz") resolved against the phase's base tick rate, which BaseSimulator._hook_base_rates() supplies from the sim config: fps for the substep phases, fps / control_decimation_steps for the frame phases. HookHandle.set_every() retunes a live hook (counter resets). Event phases (episode) accept int decimation only; CLOSE rejects any decimation.
3. Participants self-register. SimulatorBridge, VirtualGantry, and the video recorder each gain a register_hooks() that wires their existing methods to phases. The backends call participant.register_hooks(self.hooks) at setup; their step methods lose the inline calls.
4. Deterministic teardown. New BaseSimulator.close() (idempotent) emits Phase.CLOSE; train_agent.py closes the env in a finally, and DirectSimulation.cleanup() routes through env.close() / simulator.close(). Adds SimulatorBridge.close() (closes the clock-publisher socket) — it did not previously exist, and registering it for CLOSE is what surfaced the bug below.

Behavior-preserving for every shipped run: the same callbacks fire at the same rate, in the same order.

Behavior changes by user impact

Newly present (opt-in / internal; no shipped run affected)

• Phase.CLOSE teardown path. BaseSimulator.close() now emits CLOSE, running bridge.close (clock socket) and video.cleanup in reverse registration order, idempotently. Previously teardown was ad-hoc: DirectSimulation.cleanup only called video_recorder.cleanup(), and the bridge's clock socket was never explicitly closed. Both train_agent.py (new finally) and DirectSimulation now go through close().
• HookRegistry public surface on BaseSimulator.hooks. New extension point — downstream code can register per-phase callbacks (with an optional cadence) without subclassing a backend or overriding on_episode_*.
• run_sim loop emits the same phases as training. DirectSimulation.run() emits PRE_STEP/POST_STEP every physics step and FRAME_BEGIN/FRAME_END every control_decimation_steps, so hooks behave identically under run_sim and BaseTask.

Removed (internal API surface)

• BaseSimulator.on_episode_start / on_episode_end / capture_video_frame / _step_bridge deleted. These were the override seams for episode/video/bridge behavior. Callers now emit phases: base_task.py emits EPISODE_START / EPISODE_END (replacing the hasattr(simulator, "on_episode_*") guards) and the frame/step phases; sim_utils.py emits EPISODE_START at init. Any out-of-tree subclass that overrode on_episode_start/end to hook episode boundaries must migrate to simulator.hooks.add(Phase.EPISODE_START, ...). All in-tree callers migrated.
• Video recorder's manual decimation counter deleted. capture_frame no longer tracks _frame_counter % control_decimation; it registers on FRAME_END, which already fires once per frame at control rate.

Result-preserving (verified)

• Per-step callback timing unchanged. bridge.step and virtual_gantry.step register on PRE_STEP (fires every physics substep, before simulate_at_each_physics_step, as before — SDK torques and gantry forces must land before the substep integrates them). Registration order — gantry before bridge (gantry registered ahead of _init_bridge in every backend) — matches the old inline call order.
• Video capture rate unchanged. video.capture_frame registers on FRAME_END: once per frame, at control frequency — the same rate the old per-substep counter decimated to. The capture point moves from "after the Nth physics substep" to "after _refresh_sim_tensors at frame end".
• Episode-start order unchanged. Gantry's set_position_to_robot (now VirtualGantry.on_episode_start, env_id == 0 guarded) registers before the video recorder's on_episode_start, matching the old BaseSimulator.on_episode_start body.

Migration guide (out-of-tree only)

Episode / video / bridge hooks — register instead of override

# BEFORE (subclass override)                 # AFTER (register a callback)
class MySim(MuJoCo):                          sim.hooks.add(Phase.EPISODE_START, my_fn)   # my_fn(env_id)
    def on_episode_start(self, env_id=0):     sim.hooks.add(Phase.POST_STEP, my_capture)
        super().on_episode_start(env_id)      sim.hooks.add(Phase.FRAME_END, my_log, every="10Hz")
        my_fn(env_id)                         # deregister later via the returned handle:
                                              handle = sim.hooks.add(Phase.CLOSE, my_cleanup); handle.remove()

simulator.capture_video_frame() / _step_bridge() are gone — the video recorder and bridge register themselves; there is nothing to call. Payload arity is validated at emit: EPISODE_START/EPISODE_END pass env_id; all other phases pass no args.

[samuelgundry]

This is such a good cleanup. As discussed offline, I strongly prefer a hook system akin to game engines -- it's a proven design pattern. I haven't approved as I think we want to use tensor/batch env_ids now in the hook system. And a future-looking suggestion (consideration, really) to further unifiy sim backends vs. components/systems.

[samuelgundry]

Bug? capture_frame expects env_id but FRAME_END does not pass?
How can we catch these static/typing errors without eyeballs or runtime?

[samuelgundry]

I think very slight behavior change means VideoRecorder will now capture at step_count=0 (before step_count is incremented) whereas before it captured only after incrementing step_count (modulo decimation)? See lines 198 removed in video_recorder.py. So for control_decimation=4, FRAME_END now also fires on on step 0 vs. before only on step 4, 8, ... Not sure if this is intentionally?