Merge pull request #1548 from Motphys/feat/issue-21-playback-overlay-session

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

Author: TATP-233

docs/sphinx/source/adr/ADR-0000-index.md

@@ -21,6 +21,7 @@ orphan: true
 | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted |
 | [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted |
 | [ADR-0007 UniSim Extraction Boundary](ADR-0007-unisim-extraction-boundary.md) | Physics package extraction | Accepted |
+| [ADR-0008 Debug Overlay Primitive Contract And Playback Session](ADR-0008-debug-overlay-primitive-contract-and-playback-session.md) | Debug overlay / playback session | Accepted |
 
 ## ADR Governance
 

docs/sphinx/source/adr/ADR-0008-debug-overlay-primitive-contract-and-playback-session.md

@@ -0,0 +1,147 @@
+---
+orphan: true
+---
+
+# ADR-0008 Debug Overlay Primitive Contract And Embeddable Playback Session
+
+- Status: Accepted
+- Date: 2026-09-09
+- Owners: Backend / Visualization maintainers
+- Supersedes: None
+- Superseded by: None
+
+## Context
+
+播放/录制管线的任务侧叠加层过去依赖 `(num_envs, 3)` marker 位置数组
+（`extra_data_getter`），只能表达“每个 env 一个球”，无法表达坐标系、箭头、
+ghost mesh 或文本等任务 debug 语义，交互 viewer（`play_interactive.py`）与离线
+record 管线各自维护了一套互不兼容的绘制代码。
+
+上游 unisim（unilabsim/unisim#53, 驱动 issue unilabsim/wuji_unilab#21）把
+`SimBackend.run_playback` 的 `extra_data_getter` 替换为
+`debug_overlay_getter`，引入 typed `DebugPrimitive`
+（sphere/box/frame/arrow/ghost_geom/text，env 局部系位姿）、typed `CameraCfg`
+（`from_kwargs` 归一化、未知键 fail-closed）以及
+`BackendPlayCapabilities.supports_debug_overlay`。UniLab 作为下游需要：
+
+1. 把 env 契约（`ABEnv`/`NpEnv`）和所有调用点迁移到新契约；
+2. 收敛交互 viewer 的硬编码 `user_scn` 叠加层到同一原语契约；
+3. 把 record 管线的“快照缓存 + 事后渲染”暴露为可嵌入组件，供自定义 eval
+   循环（自写 trial 协议）复用。
+
+这跨越 env/backend 公共契约，普通实现说明不足以作为 review 基线。
+
+## Decision
+
+1. **叠加层一律使用 typed `DebugPrimitive`。** env 契约透传
+   `debug_overlay_getter: Callable[[], Sequence[Sequence[DebugPrimitive] | None] | None]`，
+   外层长度等于 `num_envs`，位姿为 env 局部系；grid offset 由渲染器应用。
+   旧的 `extra_data_getter` marker 数组契约在 UniLab 侧零残留。
+2. **capability 门控 fail-closed。** `EnvPlayCapabilities` 增加
+   `supports_debug_overlay` 并透传 backend capability；调用方在 backend 不
+   支持时省略 getter，backend 在收到 getter 但不支持时抛
+   `NotImplementedError`。
+3. **交互 viewer 与离线渲染共用原语数据层。** `play_interactive.py` 只构造
+   `DebugPrimitive` 列表；注入 `viewer.user_scn` 的统一入口是
+   `unilab.visualization.debug_primitives.append_debug_primitives_to_scene`。
+   该注入器是 UniLab 持有的过渡实现——unisim 的离线 worker 绘制实现
+   （`_append_primitive`/`_append_debug_primitives`）是私有的，且面向多 env
+   网格合成而非交互路径（单 env、已加载 model、调用方持有 mjvScene）；
+   待 unisim 暴露适合交互路径的公开 helper 后改为委托（见 Consequences）。
+4. **record 管线组件化为 `SnapshotPlaybackSession`。**
+   `unilab.visualization.playback_session.SnapshotPlaybackSession` 把
+   `snapshot()`（每 trial 步缓存 physics state 与可选叠加层）与
+   `render_snapshots(output_video=..., overlay_getter=..., camera=..., fps=...,
+   on_frame=...)`（统一渲染出 mp4）暴露为 session 级公开操作，渲染走共享的
+   MuJoCo 离线 snapshot 管线。session 只持有 NumPy 数组与 typed 原语，缓存
+   可 pickle；trial 之间用 `clear()` 界定生命周期。
+5. **`ABEnv.render(mode="rgb_array")` 便捷封装。** 内部走
+   `init_play_renderer(headless=True, capture=True)` +
+   `capture_play_video_frame`，按 `supports_native_video_capture` 门控，不支持
+   时抛带类名的 `NotImplementedError`。
+6. **camera 配置面统一过 `CameraCfg.from_kwargs`。** train/play 脚本通过
+   `unilab.visualization.playback.camera_cfg_from_training` 组装 typed camera
+   配置；Hydra YAML 字段名不变，未知键在边界 fail-closed。
+7. **`on_frame` 回调先在 session 层落地。** `run_playback`/`run_playback_mode`
+   在 env 契约上声明 `on_frame: Callable[[int, np.ndarray], np.ndarray | None]`，
+   但 unisim `SimBackend.run_playback` 尚未声明该参数；env 透传层在收到
+   `on_frame` 时 fail-closed 抛 `NotImplementedError`，实际帧回调由
+   `SnapshotPlaybackSession.render_snapshots(on_frame=...)` 提供，待上游补契约后
+   再透传（见 Consequences）。
+8. **任务自有 overlay 走统一发现入口。** `ManagerBasedRlEnv` 提供
+   `get_playback_debug_overlays()`：遍历 command manager 的 terms，聚合实现了
+   `playback_debug_overlay_getter()` 的 term，把多 term 的原语按 env 合并成单个
+   `DebugOverlayGetter`；无 provider 时返回 `None`。play 入口
+   （`train_rsl_rl.play_rsl_rl`、`play_interactive.py`）通过
+   `getattr(env, "get_playback_debug_overlays", None)` 发现任务 overlay，发现不到
+   时回退现有特例（`curr_ee_goal_world` EE goal sphere、交互 viewer 的
+   motion/reward/velocity 硬编码），既有行为不回归。
+9. **overlay 门控按 resolve 后的渲染模式区分。** 落点在
+   `ABEnv.run_playback_mode`（`on_plan` 回调之后、dispatch 之前，plan 对象不被
+   改写）：record 模式要求 `supports_debug_overlay`，backend 不支持时维持
+   fail-closed；interactive 模式要求
+   `supports_interactive_debug_overlay`（unisim#54 起 mjwarp 报 True，其余
+   backend 为 False），不满足时 `warnings.warn` 并丢弃 getter（传 None），
+   交互回放照常进行——任务 opt-in overlay 不应让无位 backend 的交互回放整体
+   不可用。
+
+## Stable Contracts
+
+- 任务侧叠加层的唯一数据契约是 `unisim.backend.base.DebugPrimitive` 列表
+  （per-env、env 局部系）；不允许再引入按位置数组或 backend 私有 geom 操作
+  的叠加层入口。
+- `EnvPlayCapabilities.supports_debug_overlay` 是 env 侧判断叠加层可用性的
+  唯一入口；interactive 模式的叠加层可用性由
+  `supports_interactive_debug_overlay` 表达，门控统一在 `run_playback_mode`。
+- 任务自有 overlay 的唯一发现入口是 `ManagerBasedRlEnv.get_playback_debug_overlays()`
+  （command terms 实现 `playback_debug_overlay_getter()`）；play 入口用
+  `getattr` 发现并回退现有特例，不再新增绕过发现机制的任务特例。
+- 自定义 eval 循环通过 `SnapshotPlaybackSession` 复用 record 管线，不在脚本里
+  重新实现“快照缓存 + 事后渲染”。
+- camera 参数在 train/play 入口统一经 `camera_cfg_from_training` /
+  `CameraCfg.from_kwargs` 归一化。
+
+## Alternatives Considered
+
+- 在 UniLab 复制 unisim 离线 worker 的 `_append_primitive` 实现供交互路径使用。
+  拒绝原因：跨仓库复制同一绘制实现必然漂移；交互路径只需要
+  sphere/box/frame/arrow 的 mjvScene 注入，过渡实现保持最小并明确等待上游
+  公开 helper。
+- 交互 viewer 继续使用硬编码 `user_scn` 绘制、只迁移 record 管线。拒绝原因：
+  两条路径的 overlay 语义会继续分叉，task 侧 debug 可视化无法在两种
+  渲染路径间复用。
+- 自定义 eval 循环直接调用 `env.run_playback()` 并自行包一层 trial 协议。
+  拒绝原因：`run_playback` 是单体的“initialize/step/渲染”整体入口，无法
+  表达“逐 trial 攒帧、trial 结束统一出片”的协议，会导致下游复制 record
+  管线内部逻辑。
+
+## Consequences
+
+- `extra_data_getter` 在 UniLab 全仓（含 docstring/tests）零残留；新叠加层
+  代码必须使用 `DebugPrimitive`。
+- 交互注入器 `unilab/visualization/debug_primitives.py` 是过渡组件，当
+  unisim 提供适合交互路径（单 env、已加载 model、调用方持有 mjvScene）的
+  公开 helper（建议形态：`append_debug_primitives(scene, overlays, *,
+  offsets=None, mesh_ids=None)`）后应改为委托并删除本地实现。
+- `on_frame` 透传依赖 unisim 在 `SimBackend.run_playback` 上声明同名参数；
+  上游落地前 env 层 fail-closed。
+- UniLab 依赖 unisim PR #53 的契约；该 PR 合并并发版后需 bump
+  `unisim-core` 最低版本。
+
+## Evidence In Repo
+
+- env 契约: `src/unilab/base/base.py`, `src/unilab/base/np_env.py`
+- 任务 overlay 发现: `src/unilab/envs/manager_based_rl_env.py`（`get_playback_debug_overlays`）
+- 交互注入器: `src/unilab/visualization/debug_primitives.py`
+- 可嵌入 session: `src/unilab/visualization/playback_session.py`
+- camera 归一化: `src/unilab/visualization/playback.py`
+- 交互 viewer 迁移: `src/unilab/scripts/play_interactive.py`
+- 训练入口迁移: `src/unilab/scripts/train_rsl_rl.py`, `src/unilab/scripts/train_appo.py`, `src/unilab/scripts/train_offpolicy.py`
+- 上游契约: `unisim.backend.base`（`DebugPrimitive`, `CameraCfg`, `DebugOverlayGetter`, `validate_debug_overlays`, `BackendPlayCapabilities.supports_debug_overlay` / `supports_interactive_debug_overlay`）
+- 测试: `tests/visualization/test_debug_primitives.py`, `tests/visualization/test_playback_session.py`, `tests/base/test_np_env_playback_contract.py`, `tests/envs/test_manager_based_rl_env.py`（overlay 聚合）
+
+## Related Documents
+
+- {doc}`ADR Index </adr/README>`
+- {doc}`ADR-0002 Backend Capability Boundary For Play And Snapshot </adr/ADR-0002-backend-capability-boundary-for-play-and-snapshot>`
+- {doc}`ADR-0007 UniSim Extraction Boundary </adr/ADR-0007-unisim-extraction-boundary>`

docs/sphinx/source/adr/README.md

@@ -20,6 +20,8 @@ orphan: true
 | [ADR-0004 Registry Bootstrap Contract](ADR-0004-registry-bootstrap-contract.md) | Registry bootstrap | Accepted |
 | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted |
 | [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted |
+| [ADR-0007 UniSim Extraction Boundary](ADR-0007-unisim-extraction-boundary.md) | Physics package extraction | Accepted |
+| [ADR-0008 Debug Overlay Primitive Contract And Playback Session](ADR-0008-debug-overlay-primitive-contract-and-playback-session.md) | Debug overlay / playback session | Accepted |
 
 ## ADR Governance
 

pyproject.rocm.toml

@@ -25,10 +25,10 @@ dependencies = [
     "numpy",
     # Physics implementations are provided by the independently released
     # unisim-core package from the production PyPI index.
-    "unisim-core>=1.1.4",
+    "unisim-core>=1.1.5",
     # RL algorithms and async runtimes live in the independently released
     # uni-rl package (distribution name ``unilab-rl``); see pyproject.toml.
-    "unilab-rl==1.1.1",
+    "unilab-rl==1.1.3",
     "torch==2.11.0",
     "triton-rocm==3.6.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "gymnasium",

pyproject.toml

@@ -41,12 +41,12 @@ dependencies = [
     "numpy",
     # Physics implementations are provided by the independently released
     # unisim-core package from the production PyPI index.
-    "unisim-core>=1.1.4",
+    "unisim-core>=1.1.5",
     # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners,
     # collectors, IPC, logging) live in the independently released uni-rl
     # package (distribution name ``unilab-rl``), consumed via the injected
     # env contract (uni_rl.env_contract.EnvFactory). Published on PyPI.
-    "unilab-rl==1.1.1",
+    "unilab-rl==1.1.3",
     "numba>=0.67",
     "prettytable>=3.10",
     # torch is a range (not an exact pin) so that published PyPI metadata lets

scripts/visualize_task_env.py

@@ -32,6 +32,7 @@
 from typing import TYPE_CHECKING, Any, cast, get_args, get_origin, get_type_hints
 
 import numpy as np
+from unisim.backend.base import CameraCfg
 
 ROOT_DIR = Path(__file__).parent.parent
 SRC_DIR = ROOT_DIR / "src"
@@ -165,7 +166,7 @@ def step(_obs):
     )
 
 
-def _motrix_camera_kwargs(env, num_envs: int) -> dict[str, Any] | None:
+def _motrix_camera_kwargs(env, num_envs: int) -> CameraCfg | None:
     """Point Motrix's interactive camera at the actual terrain spawn cells."""
     spawn = getattr(env, "_spawn", None)
     origins_for = getattr(spawn, "origins_for", None)
@@ -194,12 +195,14 @@ def _motrix_camera_kwargs(env, num_envs: int) -> dict[str, Any] | None:
         lookat = origins[0].copy()
         distance = 4.0
     lookat[2] += 0.5
-    return {
-        "cam_lookat": lookat.tolist(),
-        "cam_distance": distance,
-        "cam_elevation": -25.0,
-        "cam_azimuth": 135.0,
-    }
+    return CameraCfg.from_kwargs(
+        {
+            "cam_lookat": lookat.tolist(),
+            "cam_distance": distance,
+            "cam_elevation": -25.0,
+            "cam_azimuth": 135.0,
+        }
+    )
 
 
 def _env_scene(env) -> "SceneCfg | None":

src/unilab/base/base.py

@@ -1,15 +1,18 @@
 import abc
-from collections.abc import Callable
+import warnings
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from os import PathLike
 from typing import Any, Optional
 
 import gymnasium as gym
 import numpy as np
-from unisim.backend.base import BackendPlayRenderPlan
+from unisim.backend.base import BackendPlayRenderPlan, CameraCfg, DebugOverlayGetter
 
 from .scene import SceneCfg
 
+OnPlaybackFrameFn = Callable[[int, np.ndarray], "np.ndarray | None"]
+
 
 @dataclass(frozen=True)
 class EnvPlayCapabilities:
@@ -18,6 +21,8 @@ class EnvPlayCapabilities:
     supports_native_interactive_renderer: bool = False
     supports_physics_state_playback: bool = False
     supports_native_video_capture: bool = False
+    supports_debug_overlay: bool = False
+    supports_interactive_debug_overlay: bool = False
 
 
 @dataclass
@@ -251,10 +256,18 @@ def run_playback(
         headless: bool | None = None,
         record_video: bool | None = None,
         frame_state_getter: Callable[[], np.ndarray] | None = None,
-        camera_kwargs: dict[str, Any] | None = None,
-        extra_data_getter: Callable[[], np.ndarray | None] | None = None,
+        camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
+        debug_overlay_getter: DebugOverlayGetter | None = None,
+        on_frame: OnPlaybackFrameFn | None = None,
     ) -> str | None:
-        """Execute playback through the backend contract."""
+        """Execute playback through the backend contract.
+
+        ``debug_overlay_getter`` returns per-env sequences of
+        :class:`unisim.backend.base.DebugPrimitive` with env-local poses
+        (``None`` disables overlays for the frame).  ``on_frame`` is an
+        optional ``(frame_index, frame) -> frame | None`` callback applied to
+        each recorded frame before it is written to the video.
+        """
         raise NotImplementedError(f"{self.__class__.__name__} does not support playback execution")
 
     def run_playback_mode(
@@ -268,11 +281,20 @@ def run_playback_mode(
         render_spacing: float | None = None,
         render_offset_mode: str | None = None,
         frame_state_getter: Callable[[], np.ndarray] | None = None,
-        camera_kwargs: dict[str, Any] | None = None,
-        extra_data_getter: Callable[[], np.ndarray | None] | None = None,
+        camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
+        debug_overlay_getter: DebugOverlayGetter | None = None,
+        on_frame: OnPlaybackFrameFn | None = None,
         on_plan: Callable[[BackendPlayRenderPlan], None] | None = None,
     ) -> str | None:
-        """Resolve configured playback mode and execute it through the backend contract."""
+        """Resolve configured playback mode and execute it through the backend contract.
+
+        Overlay gating follows the resolved plan mode: ``record`` playback
+        requires ``play_capabilities.supports_debug_overlay`` (backends fail
+        closed otherwise); ``interactive`` playback requires
+        ``supports_interactive_debug_overlay`` — when the backend lacks it the
+        getter is dropped with a warning so interactive playback still runs
+        without overlays.
+        """
         plan = self.resolve_play_render_plan(
             play_render_mode=play_render_mode,
             play_steps=play_steps,
@@ -282,6 +304,14 @@ def run_playback_mode(
             on_plan(plan)
         if plan.mode == "none":
             return None
+        if plan.mode == "interactive" and debug_overlay_getter is not None:
+            if not self.play_capabilities.supports_interactive_debug_overlay:
+                warnings.warn(
+                    f"{self.__class__.__name__} backend does not support interactive debug "
+                    "overlays; dropping debug_overlay_getter for this interactive playback",
+                    stacklevel=2,
+                )
+                debug_overlay_getter = None
         return self.run_playback(
             initialize=initialize,
             step=step,
@@ -293,7 +323,8 @@ def run_playback_mode(
             record_video=plan.record_video,
             frame_state_getter=frame_state_getter,
             camera_kwargs=camera_kwargs,
-            extra_data_getter=extra_data_getter,
+            debug_overlay_getter=debug_overlay_getter,
+            on_frame=on_frame,
         )
 
     @property
@@ -351,13 +382,24 @@ def init_play_renderer(
         capture: bool = False,
         width: int = 1280,
         height: int = 720,
-        camera_kwargs: dict[str, Any] | None = None,
+        camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
     ) -> None:
         """Initialize env-facing playback rendering when supported."""
         raise NotImplementedError(
             f"{self.__class__.__name__} does not support native playback rendering"
         )
 
+    def render(self, mode: str = "rgb_array") -> np.ndarray:
+        """Render the current state; ``mode="rgb_array"`` returns an RGB frame.
+
+        Convenience wrapper over :meth:`init_play_renderer` and
+        :meth:`capture_play_video_frame`, gated on
+        ``play_capabilities.supports_native_video_capture``.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support render(mode={mode!r})"
+        )
+
     def render_play_frame(self) -> None:
         """Render one frame through the env-facing interactive playback contract."""
         raise NotImplementedError(