feat(playback): adopt unisim debug overlay primitives and embeddable playback session - #1548
Merged
Merged
Conversation
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.
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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Migrate UniLab's playback path to the new unisim debug overlay primitive contract (unilabsim/unisim#53, driving issue unilabsim/wuji_unilab#21):
base/base.py,base/np_env.py):run_playback/run_playback_modenow takedebug_overlay_getter: DebugOverlayGetter | Noneinstead ofextra_data_getter;camera_kwargsacceptsCameraCfg | Mapping[str, Any] | NonewithCameraCfg.from_kwargsnormalization at the boundary;EnvPlayCapabilitiesforwardssupports_debug_overlay. Theon_framehook is declared in the env contract but fails closed withNotImplementedErrorinNpEnv.run_playbackuntil unisim exposes it onSimBackend.run_playback(see "Upstream gaps").ABEnv.render(mode="rgb_array")is added with a lazy play-renderer capture path and capability gating.envs/manager_based_rl_env.py):ManagerBasedRlEnv.get_playback_debug_overlays()collects command terms implementingplayback_debug_overlay_getter()and aggregates them into a singleDebugOverlayGetter, merging primitives per env across terms (Nonewhen no term opts in). Play entrypoints discover it viagetattr(env, "get_playback_debug_overlays", None):train_rsl_rl.play_rsl_rlprefers task-owned overlays and falls back to thecurr_ee_goal_worldEE-goal sphere special case;play_interactive.pyapplies discovered overlays for the viewed env alongside its existing motion/reward/velocity builders. Behavior is unchanged when no term provides an overlay.base/base.py,base/np_env.py):EnvPlayCapabilitiesforwards the newsupports_interactive_debug_overlay(unisim#54; mjwarp reportsTrue, other backendsFalse).ABEnv.run_playback_modegatesdebug_overlay_getteron the resolved plan mode —recordkeeps requiringsupports_debug_overlay(backends fail closed), whileinteractiveplayback on a backend without interactive overlay support drops the getter with awarnings.warnso the session still runs without overlays. Theon_plancallback fires before gating and the plan object is not mutated.train_rsl_rl.pyrenders the end-effector goal as aDebugPrimitivesphere overlay (env-local, per-envNoneon non-finite values);train_appo.py,train_offpolicy.py,play_hora_appo.pyuse the newcamera_cfg_from_training()helper;train_hora_distill.pyandvisualize_task_env.pybuildCameraCfg.from_kwargs(...).play_interactive.py): marker drawing converges ontoDebugPrimitivebuilders (motion-target frames, reward debug vectors, velocity-command arrows) injected via the newunilab.visualization.debug_primitivesmodule (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").visualization/playback_session.py): newSnapshotPlaybackSession—snapshot()captures detached physics-state copies,render_snapshots()renders them through the offline pipeline with per-calloverlay_getteroverride, camera/fps control,on_framepost-processing, andvalidate_debug_overlayschecks, all gated by backend play capabilities.Backend impact
debug_overlay_getter(per unisim#53 capabilities).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.NotImplementedError) when a getter is passed andsupports_debug_overlayisFalse.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_indexfails 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 intests/scripts/test_train_scripts.pypass 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 onmainwith the locked unisim-core 1.1.4:tests/base/test_mjwarp_*.py—mujoco-warpoptional extra not installed in this environment (MjwarpDependencyError); these test files are untouched by this PR.tests/envs/test_sharpa.pynative-crash (segfault/abort) insidemujoco_unibatch reset; reproduces identically onmainwith unisim-core 1.1.4.make test-benchmark-smoke: PASS (mlx cases skipped as platform-optional)pytest tests/scripts/test_check_docs.pyPASS; sphinx HTML build (-n, autodoc skipped) PASStests/base/test_np_env_playback_contract.py(passthrough,on_framefail-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 intests/envs/test_manager_based_rl_env.py, plus discovery/fallback cases appended totests/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 lockedunisim-core==1.1.4does not carry the new contract, so CI cannot go green until those PRs merge and a release is published; theunisim-core>=1.1.4floor inpyproject.tomlmust then be bumped before this PR merges.pyproject.tomlis intentionally left unchanged here.Upstream gaps (requests for unisim)
unisim.visualization.render_manyonly exposes private_append_primitive/_append_debug_primitivesaimed at the multi-env worker grid. Please publish something likeunisim.visualization.render_many.append_debug_primitives(scene, overlays, *, offsets=None, mesh_ids=None) -> intfor the single-env, already-loaded-model, caller-ownedviewer.user_scncase, so UniLab can drop its transitionalunilab/visualization/debug_primitives.py.on_frameonSimBackend.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.pyis a self-contained copy of the old marker API (noextra_data_getterreferences; onlyscripts/play_viser.pyuses itsget_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 ofscripts/train_hora_distill.pyandsrc/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 inunilabsim/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.