unilab.visualization.playback_session

Embeddable snapshot-cache + deferred-render playback session.

Custom eval loops with their own trial protocol cannot use the monolithic env.run_playback() entrypoint. SnapshotPlaybackSession exposes the record pipeline’s two halves as session-level operations: call SnapshotPlaybackSession.snapshot() once per trial step to cache physics states (and, optionally, per-frame debug overlays), then call SnapshotPlaybackSession.render_snapshots() at the end of a trial to render the cached states into one mp4 through the shared MuJoCo offline snapshot pipeline.

The session holds plain NumPy arrays and typed primitives only, so cached state stays picklable; the owning env reference is used for capability gating and model resolution at render time. Call SnapshotPlaybackSession.clear() between trials to bound memory.

Classes

SnapshotPlaybackSession

Cache physics snapshots during a trial; render them to video afterwards.
class unilab.visualization.playback_session.SnapshotPlaybackSession[source]

Bases: object

Cache physics snapshots during a trial; render them to video afterwards.

Parameters:
__init__(env, *, frame_state_getter=None, overlay_getter=None, render_spacing=None, width=1280, height=720, num_processes=8)[source]
Parameters:
property snapshots: tuple[ndarray, ...]

The cached physics states, in capture order.

snapshot()[source]

Cache one physics snapshot (and the current overlays, if configured).

Return type:

ndarray

clear()[source]

Drop all cached snapshots and overlays (call between trials).

Return type:

None

render_snapshots(*, output_video, overlay_getter=None, camera=None, fps=None, on_frame=None)[source]

Render cached snapshots to output_video and return its path.

overlay_getter here is evaluated once per cached frame at render time and overrides the snapshot-time overlays captured through the constructor getter; prefer the constructor overlay_getter for state-coupled overlays. camera is normalized through CameraCfg.from_kwargs() (unknown keys fail closed). on_frame receives (frame_index, frame) after rendering and may return a modified frame. Returns None when the host cannot render off-screen (a warning is printed by the renderer).

Parameters:
Return type:

str | None