Skip to content

feat(playback): adopt unisim debug overlay primitives and embeddable playback session - #1548

Merged
TATP-233 merged 8 commits into
mainfrom
feat/issue-21-playback-overlay-session
Sep 9, 2026
Merged

feat(playback): adopt unisim debug overlay primitives and embeddable playback session#1548
TATP-233 merged 8 commits into
mainfrom
feat/issue-21-playback-overlay-session

Conversation

@TATP-233

@TATP-233 TATP-233 commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

Migrate UniLab's playback path to the new unisim debug overlay primitive contract (unilabsim/unisim#53, driving issue unilabsim/wuji_unilab#21):

  • Env contract (base/base.py, base/np_env.py): run_playback / run_playback_mode now take debug_overlay_getter: DebugOverlayGetter | None instead of extra_data_getter; camera_kwargs accepts CameraCfg | Mapping[str, Any] | None with CameraCfg.from_kwargs normalization at the boundary; EnvPlayCapabilities forwards supports_debug_overlay. The on_frame hook is declared in the env contract but fails closed with NotImplementedError in NpEnv.run_playback until unisim exposes it on SimBackend.run_playback (see "Upstream gaps"). ABEnv.render(mode="rgb_array") is added with a lazy play-renderer capture path and capability gating.
  • Task-owned overlay discovery (envs/manager_based_rl_env.py): ManagerBasedRlEnv.get_playback_debug_overlays() collects command terms implementing playback_debug_overlay_getter() and aggregates them into a single DebugOverlayGetter, merging primitives per env across terms (None when no term opts in). Play entrypoints discover it via getattr(env, "get_playback_debug_overlays", None): train_rsl_rl.play_rsl_rl prefers task-owned overlays and falls back to the curr_ee_goal_world EE-goal sphere special case; play_interactive.py applies discovered overlays for the viewed env alongside its existing motion/reward/velocity builders. Behavior is unchanged when no term provides an overlay.
  • Mode-aware overlay gating (base/base.py, base/np_env.py): EnvPlayCapabilities forwards the new supports_interactive_debug_overlay (unisim#54; mjwarp reports True, other backends False). ABEnv.run_playback_mode gates debug_overlay_getter on the resolved plan mode — record keeps requiring supports_debug_overlay (backends fail closed), while interactive playback on a backend without interactive overlay support drops the getter with a warnings.warn so the session still runs without overlays. The on_plan callback fires before gating and the plan object is not mutated.
  • Call sites: train_rsl_rl.py renders the end-effector goal as a DebugPrimitive sphere overlay (env-local, per-env None on non-finite values); train_appo.py, train_offpolicy.py, play_hora_appo.py use the new camera_cfg_from_training() helper; train_hora_distill.py and visualize_task_env.py build CameraCfg.from_kwargs(...).
  • Interactive viewer (play_interactive.py): marker drawing converges onto DebugPrimitive builders (motion-target frames, reward debug vectors, velocity-command arrows) injected via the new unilab.visualization.debug_primitives module (quat_from_z_axis, segment_arrow, append_debug_primitives_to_scene). This injector is a transitional UniLab-held helper until unisim publishes a public interactive-path API (see "Upstream gaps").
  • Embeddable playback (visualization/playback_session.py): new SnapshotPlaybackSessionsnapshot() captures detached physics-state copies, render_snapshots() renders them through the offline pipeline with per-call overlay_getter override, camera/fps control, on_frame post-processing, and validate_debug_overlays checks, all gated by backend play capabilities.
  • ADR: ADR-0008 records the contract migration, the session API, the task-overlay discovery hook, and the transitional injector (index and README updated; the previously missing ADR-0007 README entry is added too).

Backend impact

  • MuJoCo / mjwarp / drake / newton / superdex offline playback pipelines support debug_overlay_getter (per unisim#53 capabilities).
  • Interactive overlay injection is supported by mjwarp only (supports_interactive_debug_overlay=True, unisim#54); interactive playback on other backends drops the getter with a warning instead of failing closed, so task opt-in overlays never block plain interactive playback.
  • All other backends fail closed (NotImplementedError) when a getter is passed and supports_debug_overlay is False.

Validation

  • make check: PASS (mypy 142 files 0 errors, pyright 0 errors, ruff clean)
  • make test: 1636 passed, 1 failed — tests/base/test_backend_imports.py::test_unisim_dependency_is_installed_from_package_index fails only because this dev environment installs unisim as a local editable (required to develop against the unisim#53/fix:remove motrix legacy #54 branch); it passes with the index package in CI.
  • make test-slow: 831 passed, 25 skipped (full suite, excluding the sharpa crashers below); the 4 new discovery tests in tests/scripts/test_train_scripts.py pass in a focused slow run at the final head (164 passed in that file). Remaining failures are pre-existing environment issues, verified unrelated by reproducing on main with the locked unisim-core 1.1.4:
    • 18 failed + 10 errors in tests/base/test_mjwarp_*.pymujoco-warp optional extra not installed in this environment (MjwarpDependencyError); these test files are untouched by this PR.
    • 5 tests in tests/envs/test_sharpa.py native-crash (segfault/abort) inside mujoco_uni batch reset; reproduces identically on main with unisim-core 1.1.4.
  • make test-benchmark-smoke: PASS (mlx cases skipped as platform-optional)
  • Docs: pytest tests/scripts/test_check_docs.py PASS; sphinx HTML build (-n, autodoc skipped) PASS
  • New tests: tests/base/test_np_env_playback_contract.py (passthrough, on_frame fail-closed, interactive overlay gating: drop-with-warning / forward / record untouched), tests/visualization/test_debug_primitives.py, tests/visualization/test_playback_session.py, overlay aggregation cases in tests/envs/test_manager_based_rl_env.py, plus discovery/fallback cases appended to tests/scripts/test_train_scripts.py.

Dependency / merge order

This PR requires unilabsim/unisim#53 and #54 (developed against local unisim branch feat/issue-21-debug-overlay-primitives @ 0168a27, package version 1.1.5). The locked unisim-core==1.1.4 does not carry the new contract, so CI cannot go green until those PRs merge and a release is published; the unisim-core>=1.1.4 floor in pyproject.toml must then be bumped before this PR merges. pyproject.toml is intentionally left unchanged here.

Upstream gaps (requests for unisim)

  1. Public interactive-path overlay helper. unisim.visualization.render_many only exposes private _append_primitive / _append_debug_primitives aimed at the multi-env worker grid. Please publish something like unisim.visualization.render_many.append_debug_primitives(scene, overlays, *, offsets=None, mesh_ids=None) -> int for the single-env, already-loaded-model, caller-owned viewer.user_scn case, so UniLab can drop its transitional unilab/visualization/debug_primitives.py.
  2. on_frame on SimBackend.run_playback, e.g. on_frame: Callable[[int, np.ndarray], np.ndarray | None] = None, so the per-frame hook declared in UniLab's env contract can be honored by backend-owned playback instead of failing closed.

Leftovers

  • src/unilab/visualization/render_many.py is a self-contained copy of the old marker API (no extra_data_getter references; only scripts/play_viser.py uses its get_grid_offsets; tests depend on its internals) — untouched here, candidate for a follow-up cleanup.

Refs unilabsim/wuji_unilab#21
ADR: docs/sphinx/source/adr/ADR-0008-debug-overlay-primitive-contract-and-playback-session.md

Merge conflict resolution

Merged origin/main (incl. PR #1549 removing the sharpa task and hora algorithm) at 6f2bfe3: accepted the upstream deletion of scripts/train_hora_distill.py and src/unilab/scripts/play_hora_appo.py (this branch's CameraCfg/overlay edits on them are superseded by the extraction; the equivalent migration is being re-applied in unilabsim/sharpa_rl_unilab), and removed this branch's hora-camera test plus the ADR-0008 evidence entry for the deleted file. No residual hora/sharpa references remain on this branch.

Replace extra_data_getter with debug_overlay_getter across
ABEnv.run_playback/run_playback_mode and NpEnv, declare the on_frame
hook in the env contract with a fail-closed NotImplementedError until
unisim exposes it on SimBackend.run_playback, and accept
CameraCfg | Mapping for camera_kwargs with CameraCfg.from_kwargs
normalization at the boundary.

Add EnvPlayCapabilities.supports_debug_overlay passthrough and
ABEnv.render(mode="rgb_array") backed by the lazy play-renderer
capture path, migrate play call sites to camera_cfg_from_training /
CameraCfg.from_kwargs, and render the rsl_rl end-effector goal as a
DebugPrimitive sphere overlay.

Refs unilabsim/wuji_unilab#21
…ack session

Add unilab.visualization.debug_primitives with quat_from_z_axis,
segment_arrow, and append_debug_primitives_to_scene as a transitional
UniLab-held injector for the interactive viewer path until unisim
publishes a public helper, and converge play_interactive.py marker
drawing onto DebugPrimitive builders (motion target frames, reward
debug vectors, velocity command arrows).

Add SnapshotPlaybackSession for embeddable playback: snapshot()
captures detached physics-state copies, and render_snapshots() renders
them through the offline pipeline with per-call overlay_getter
override, camera/fps control, on_frame frame post-processing, and
validate_debug_overlays checks, gated by backend play capabilities.

Refs unilabsim/wuji_unilab#21
Document the migration to the unisim debug overlay primitive contract,
the SnapshotPlaybackSession embedding API, the fail-closed on_frame
stance, and the transitional UniLab-held scene injector pending a
public unisim helper. Also register the missing ADR-0007 entry in the
ADR README index.

Refs unilabsim/wuji_unilab#21
Add ManagerBasedRlEnv.get_playback_debug_overlays(), which collects
command terms implementing playback_debug_overlay_getter() and
aggregates them into a single DebugOverlayGetter, merging primitives
per env across terms and returning None when no term provides one.

train_rsl_rl.play_rsl_rl now discovers task-owned overlays via
getattr(env, "get_playback_debug_overlays", None) through the new
_playback_debug_overlay_getter, falling back to the curr_ee_goal_world
EE-goal special case; play_interactive.py applies discovered overlays
for the viewed env alongside its existing motion/reward/velocity
builders, so behavior is unchanged when no term opts in.

Refs unilabsim/wuji_unilab#21
BackendPlayCapabilities now distinguishes supports_interactive_debug_overlay
(unisim#54: mjwarp reports True, other backends False). Forward it through
EnvPlayCapabilities and gate debug_overlay_getter in
ABEnv.run_playback_mode on the resolved plan mode: record keeps requiring
supports_debug_overlay (backends fail closed), while interactive playback
on a backend without interactive overlay support drops the getter with a
warning so the session still runs without overlays. The on_plan callback
fires before gating and the plan object is not mutated.

Refs unilabsim/wuji_unilab#21
…branch

Accept the upstream deletion of scripts/train_hora_distill.py and
src/unilab/scripts/play_hora_appo.py (PR #1549 moved hora to
unilabsim/sharpa_rl_unilab); the CameraCfg/overlay edits on those files
are superseded by the move and will be re-applied in the new repo.
Drop the hora camera test appended by this branch and the ADR-0008
evidence entry for play_hora_appo.py.

Refs unilabsim/wuji_unilab#21
The mjwarp snapshot layout appends [mocap_pos, mocap_quat] after
[time, qpos, qvel] for models with mocap bodies (unisim playback fix for
record-mode mocap interpenetration). Extend the backend stub with the
mocap host arrays, relax the layout error regex for the extended layout
message, and add a mocap snapshot test.
@TATP-233

TATP-233 commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator Author

Dependencies bumped to the released versions: unisim-core >=1.1.5, unilab-rl ==1.1.3 (pyproject.toml + pyproject.rocm.toml, both locks regenerated via uv lock). Local gates on 8d00117: make check clean, make test 1573 passed, make test-all passed (34/35 module-mode, 35/36 script-mode, mlx-only skips). The pyright / benchmark-smoke / ubuntu-slim failures on the previous head were caused by unisim-core 1.1.5 not yet being on PyPI; this head consumes the published release.

@TATP-233
TATP-233 merged commit 45e4dff into main Sep 9, 2026
13 of 14 checks passed
@TATP-233
TATP-233 deleted the feat/issue-21-playback-overlay-session branch September 9, 2026 16:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant