"""Playback rendering entrypoint for trained policies."""
from __future__ import annotations
from pathlib import Path
from typing import Any, Callable, TypeVar, cast
import numpy as np
import torch
from unisim.backend.base import CameraCfg, DebugOverlayGetter
ObsT = TypeVar("ObsT")
[docs]
def camera_cfg_from_training(training_cfg: Any) -> CameraCfg:
"""Assemble the typed playback camera config from a training owner config.
Hydra ``training.cam_*`` fields keep their names; this helper is the single
normalization point so unknown keys fail closed at the unisim boundary
instead of being silently ignored.
"""
return CameraCfg.from_kwargs(
{
"cam_distance": training_cfg.cam_distance,
"cam_elevation": training_cfg.cam_elevation,
"cam_azimuth": training_cfg.cam_azimuth,
"cam_lookat": getattr(training_cfg, "cam_lookat", None),
"cam_tracking": getattr(training_cfg, "cam_tracking", False),
"cam_tracking_env_idx": getattr(training_cfg, "cam_tracking_env_idx", 0),
"cam_tracking_extra_envs": getattr(training_cfg, "cam_tracking_extra_envs", 2),
"cam_fov": getattr(training_cfg, "cam_fov", None),
}
)
[docs]
def render_play_mode(
env,
*,
initialize: Callable[[], ObsT],
step: Callable[[ObsT], ObsT],
num_steps: int | None,
output_video: str | Path | None = None,
render_spacing: float | None = None,
render_offset_mode: str | None = None,
headless: bool | None = None,
record_video: bool | None = None,
frame_state_getter: Callable[[], np.ndarray] | None = None,
camera_kwargs: CameraCfg | dict[str, Any] | None = None,
debug_overlay_getter: DebugOverlayGetter | None = None,
on_frame: Callable[[int, np.ndarray], np.ndarray | None] | None = None,
) -> str | None:
"""Run playback through the env/backend playback contract.
Backend selection belongs to ``env.run_playback()`` and the concrete
backend implementations.
"""
return cast(
str | None,
env.run_playback(
initialize=initialize,
step=step,
num_steps=num_steps,
output_video=output_video,
render_spacing=render_spacing,
render_offset_mode=render_offset_mode,
headless=headless,
record_video=record_video,
frame_state_getter=frame_state_getter,
camera_kwargs=camera_kwargs,
debug_overlay_getter=debug_overlay_getter,
on_frame=on_frame,
),
)
[docs]
def run_native_viewer_play_loop(
env,
*,
initialize: Callable[[], ObsT],
step: Callable[[ObsT], ObsT],
num_steps: int | None = None,
render_spacing: float | None = None,
render_offset_mode: str | None = None,
) -> None:
"""Run the backend-native viewer loop under ``torch.inference_mode()``.
Shared by the training entrypoints' native-viewer playback paths; callers
own the policy wiring through the ``initialize``/``step`` callbacks.
"""
with torch.inference_mode():
env.run_playback(
render_spacing=render_spacing,
render_offset_mode=render_offset_mode,
num_steps=num_steps,
initialize=initialize,
step=step,
)
__all__ = [
"camera_cfg_from_training",
"render_play_mode",
"run_native_viewer_play_loop",
]