"""Shared core for interactive policy playback entrypoints."""
from __future__ import annotations
import copy
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, ClassVar, Protocol, cast
import numpy as np
import torch
from omegaconf import DictConfig, OmegaConf
from unilab.utils.checkpoint import normalize_checkpoint_value
from unilab.utils.sim2sim import policy_load_dim_guard, resolve_sim2sim_config
LogFn = Callable[[str], None]
[docs]
@dataclass(frozen=True)
class PlaybackConfig:
"""Configuration needed to bootstrap an interactive playback session."""
task: str
load_run: str
checkpoint: str | None
action_mode: str
policy_obs_mode: str
algo_log_name: str
log_root: str | None
num_envs: int = 1
speed: float = 1.0
start_paused: bool = False
[docs]
@dataclass
class PlayInteractiveArgs:
"""Scalar play arguments shared by interactive playback entrypoints."""
task: str
load_run: str
checkpoint: str | None
action_mode: str
policy_obs_mode: str
algo_log_name: str
log_root: str | None
show_target_bodies: bool
show_reward_debug: bool
target_show_axes: bool
target_body_names: str
target_max_bodies: int
target_marker_radius: float
target_axis_length: float
target_marker_alpha: float
reward_debug_show_velocity: bool
reward_debug_lin_vel_scale: float
reward_debug_ang_vel_scale: float
reward_debug_show_connectors: bool
reward_debug_show_global_anchor: bool
camera_follow_body: bool
camera_focus_body_name: str
camera_height_offset: float
camera_distance: float | None
camera_elevation: float | None
camera_azimuth: float | None
use_env_visual_model: bool
speed: float
start_paused: bool
keyboard: bool = False
keyboard_step_lin: float = 0.1
keyboard_step_ang: float = 0.2
require_keyboard_command_obs: bool = True
algo: str = "ppo"
# Physics owner selected by the interactive CLI. MuJoCo remains the viewer
# implementation for every value, including ``drake``.
sim: str = "mujoco"
[docs]
def build_playback_config(args: Any, *, num_envs: int = 1) -> PlaybackConfig:
"""Build an :class:`PlaybackConfig` from a play-args-like object."""
return PlaybackConfig(
task=str(args.task),
load_run=str(args.load_run),
checkpoint=getattr(args, "checkpoint", None),
action_mode=str(args.action_mode),
policy_obs_mode=str(args.policy_obs_mode),
algo_log_name=str(getattr(args, "algo_log_name", "rsl_rl_ppo")),
log_root=getattr(args, "log_root", None),
num_envs=num_envs,
speed=float(getattr(args, "speed", 1.0)),
start_paused=bool(getattr(args, "start_paused", False)),
)
[docs]
@dataclass
class PlaybackControls:
"""Viewer-independent playback control state."""
paused: bool = False
speed: float = 1.0
_single_step_requests: int = field(default=0, init=False, repr=False)
[docs]
def pause(self) -> None:
self.paused = True
[docs]
def resume(self) -> None:
self.paused = False
[docs]
def toggle_pause(self) -> bool:
self.paused = not self.paused
return self.paused
[docs]
def request_single_step(self, count: int = 1) -> None:
self._single_step_requests += max(int(count), 0)
[docs]
def set_speed(self, value: float) -> None:
self.speed = max(float(value), 1e-6)
[docs]
def consume_step_permission(self) -> bool:
if self.paused:
if self._single_step_requests <= 0:
return False
self._single_step_requests -= 1
return True
if self._single_step_requests > 0:
self._single_step_requests -= 1
return True
[docs]
def target_dt(self, ctrl_dt: float) -> float:
return float(ctrl_dt) / max(float(self.speed), 1e-6)
[docs]
@dataclass
class KeyboardCommander:
"""Mutable ``[vx, vy, vyaw]`` velocity command driven by keyboard nudges.
Per-axis nudges stack and are clamped to the task's ``commands.vel_limit``.
"""
low: np.ndarray
high: np.ndarray
step_lin: float = 0.1
step_ang: float = 0.2
command: np.ndarray = field(init=False)
AXIS_VX: ClassVar[int] = 0
AXIS_VY: ClassVar[int] = 1
AXIS_VYAW: ClassVar[int] = 2
def __post_init__(self) -> None:
self.low = np.asarray(self.low, dtype=np.float64).reshape(3)
self.high = np.asarray(self.high, dtype=np.float64).reshape(3)
self.command = np.zeros(3, dtype=np.float64)
[docs]
@classmethod
def from_vel_limit(
cls, vel_limit: Any, *, step_lin: float = 0.1, step_ang: float = 0.2
) -> "KeyboardCommander":
limit = np.asarray(vel_limit, dtype=np.float64)
if limit.shape != (2, 3):
raise ValueError(f"commands.vel_limit must have shape (2, 3), got {limit.shape}")
return cls(low=limit[0], high=limit[1], step_lin=float(step_lin), step_ang=float(step_ang))
[docs]
def nudge(self, axis: int, sign: float) -> None:
base = self.step_lin if axis in (self.AXIS_VX, self.AXIS_VY) else self.step_ang
delta = base * (1.0 if sign >= 0 else -1.0)
self.command[axis] = float(
np.clip(self.command[axis] + delta, self.low[axis], self.high[axis])
)
[docs]
def zero(self) -> None:
self.command[:] = 0.0
[docs]
def describe(self) -> str:
return (
f"cmd vx={self.command[0]:+.2f} vy={self.command[1]:+.2f} vyaw={self.command[2]:+.2f}"
)
[docs]
@dataclass(frozen=True)
class MotionOverlaySelection:
"""Cold-path selection of task bodies used by playback overlays."""
enabled: bool
selected_indices: np.ndarray
[docs]
class PlaybackSession(Protocol):
"""Viewer-facing session contract shared by all policy families."""
env: Any
[docs]
def reset(self) -> Any: ...
[docs]
def advance(self, controls: PlaybackControls) -> bool: ...
[docs]
def physics_state(self) -> np.ndarray: ...
@property
def info(self) -> dict[str, Any]: ...
[docs]
class RslRlPlaybackSession:
"""Policy/action stepping core shared by native and web viewers."""
[docs]
def __init__(
self,
*,
env: Any,
wrapped_env: Any,
device: str,
action_mode: str,
policy: Callable[[Any], Any] | None,
num_envs: int,
runner: Any | None = None,
actor: Any | None = None,
) -> None:
self.env = env
self.wrapped_env = wrapped_env
self.device = device
self.action_mode = action_mode
self.policy = policy
self.num_envs = int(num_envs)
self.runner = runner
self.actor = actor
self.obs: Any | None = None
self.step_count = 0
[docs]
def reset(self) -> Any:
self.obs, _info = self.wrapped_env.reset()
self.step_count = 0
return self.obs
[docs]
def step_once(self) -> Any:
actions = self._build_actions()
self.obs, _reward, _done, _info = self.wrapped_env.step(actions)
self.step_count += 1
return self.obs
[docs]
def advance(self, controls: PlaybackControls) -> bool:
if not controls.consume_step_permission():
return False
self.step_once()
return True
[docs]
def physics_state(self) -> np.ndarray:
return self.env.get_physics_state_snapshot()
@property
def info(self) -> dict[str, Any]:
state = getattr(self.env, "state", None)
info = getattr(state, "info", None)
return info if isinstance(info, dict) else {}
def _build_actions(self) -> torch.Tensor:
if self.obs is None:
raise RuntimeError("Playback session must be reset before stepping.")
action_space = self.env.action_space
action_dim = int(action_space.shape[0])
if self.action_mode == "policy" and self.policy is not None:
# Backends may expose double-precision observations (Drake's
# native state/sensor path does). Policies are trained and stored
# in float32, so normalize the playback tensor at this boundary.
if isinstance(self.obs, torch.Tensor) and self.obs.dtype != torch.float32:
self.obs = self.obs.float()
return self.policy(self.obs)
if self.action_mode == "random":
low = torch.as_tensor(action_space.low, device=self.device).to(torch.float32)
high = torch.as_tensor(action_space.high, device=self.device).to(torch.float32)
return low + (high - low) * torch.rand(
(self.num_envs, action_dim), device=self.device, dtype=torch.float32
)
return torch.zeros(self.num_envs, action_dim, device=self.device)
[docs]
class OffPolicyPlaybackSession:
"""Direct env stepping session for SAC-style off-policy actors."""
[docs]
def __init__(
self,
*,
env: Any,
device: str,
action_mode: str,
actor: Any | None,
actor_algo_type: str,
normalizer: Any | None,
num_envs: int,
obs_extractor: Callable[[dict[str, torch.Tensor]], torch.Tensor],
) -> None:
self.env = env
# A backend may remap a generic CUDA request to a concrete device.
# ``env.device`` is authoritative once the environment exists.
self.policy_device = (
torch.device(device)
if getattr(env, "device", None) is None
else torch.device(env.device)
)
self.action_mode = action_mode
self.actor = actor
self.actor_algo_type = str(actor_algo_type)
self.normalizer = normalizer
self.num_envs = int(num_envs)
self.obs_extractor = obs_extractor
self.obs: torch.Tensor | None = None
self.step_count = 0
[docs]
def reset(self) -> torch.Tensor:
if self.env.state is None:
self.env.init_state()
reset_result = self.env.reset()
if not isinstance(reset_result, tuple) or len(reset_result) != 2:
raise ValueError(f"Unexpected env.reset return format: {type(reset_result)!r}")
obs_out, _ = reset_result
self.obs = _to_play_tensor(self.obs_extractor(obs_out), device=self.policy_device)
self.step_count = 0
return self.obs
[docs]
def step_once(self) -> torch.Tensor:
actions = self._build_actions()
state = self.env.step(actions)
self.obs = _to_play_tensor(self.obs_extractor(state.obs), device=self.policy_device)
self.step_count += 1
return self.obs
[docs]
def advance(self, controls: PlaybackControls) -> bool:
if not controls.consume_step_permission():
return False
self.step_once()
return True
[docs]
def physics_state(self) -> np.ndarray:
return self.env.get_physics_state_snapshot()
@property
def info(self) -> dict[str, Any]:
state = getattr(self.env, "state", None)
info = getattr(state, "info", None)
return info if isinstance(info, dict) else {}
def _build_actions(self) -> torch.Tensor:
if self.obs is None:
raise RuntimeError("Playback session must be reset before stepping.")
action_space = self.env.action_space
action_dim = int(action_space.shape[0])
if self.action_mode == "policy" and self.actor is not None:
obs_torch = self.obs.to(device=self.policy_device, dtype=torch.float32)
if self.normalizer is not None:
obs_torch = self.normalizer(obs_torch, update=False)
actions = self.actor.explore(obs_torch, deterministic=True)
return actions.detach().to(device=self.env.device, dtype=torch.float32).contiguous()
if self.action_mode == "random":
low = torch.as_tensor(action_space.low, device=self.env.device).to(torch.float32)
high = torch.as_tensor(action_space.high, device=self.env.device).to(torch.float32)
return low + (high - low) * torch.rand(
(self.num_envs, action_dim), device=self.env.device, dtype=torch.float32
)
return torch.zeros((self.num_envs, action_dim), dtype=torch.float32, device=self.env.device)
def _to_play_tensor(value: Any, *, device: torch.device) -> torch.Tensor:
if not isinstance(value, torch.Tensor):
raise TypeError(
"Off-policy playback observations must be torch tensors; "
f"received {type(value).__name__}"
)
return value.to(device=device, dtype=torch.float32).contiguous()
[docs]
def make_sim2sim_preflight(
cfg: DictConfig,
*,
algo_name: str,
) -> Callable[[str | None], Any]:
"""Build a sim2sim contract preflight closure for interactive playback entrypoints.
Requires the composed Hydra config for the playback run; the legacy
non-Hydra path (which silently skipped the preflight) has been removed, so
a missing or non-Hydra config fails closed. A source run whose contract
snapshot is missing, malformed, or uses retired keys also fails closed in
:func:`resolve_sim2sim_config`; only an explicit ``None`` source run dir
skips the check.
"""
if not isinstance(cfg, DictConfig):
raise ValueError(
"make_sim2sim_preflight requires a composed Hydra config (DictConfig); "
"interactive playback no longer supports the legacy non-Hydra path and "
"never skips the sim2sim preflight. Compose the owner config with "
f"hydra.compose and pass it to the playback entrypoint; got "
f"{type(cfg).__name__}."
)
def _preflight(source_run_dir: str | None) -> Any:
return resolve_sim2sim_config(
source_run_dir,
cfg,
algo_name=algo_name,
strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
)
return _preflight
def _scene_visual_materializer() -> Any:
from unilab.base import config_adapter
return config_adapter.materialize_scene_visual_override
[docs]
def build_play_backend_adapter(cfg: Any, *, root_dir: str | Path, algo_name: str = "ppo") -> Any:
"""Build the EnvConfigAdapter used by play entrypoints to derive env cfg overrides."""
from unilab.base.config_adapter import EnvConfigAdapter
return EnvConfigAdapter(
cfg,
root_dir=root_dir,
algo_name=algo_name,
scene_materializer=_scene_visual_materializer(),
)
[docs]
def available_backends_for_task(task_name: str) -> tuple[str, ...]:
"""Return the registered backends for a task, or ``()`` when unknown."""
from unilab.base import registry
envs = registry.list_registered_envs()
task_meta = envs.get(task_name, {})
backends = task_meta.get("available_backends", ())
if not isinstance(backends, list):
return ()
return tuple(str(backend) for backend in backends)
[docs]
def create_rsl_rl_playback_session(
*,
playback_cfg: PlaybackConfig,
env_factory: Callable[[int], Any],
algo_config: dict[str, Any],
root_dir: str | Path,
device: str | None,
checkpoint_resolver: Callable[[str, str, str | None, str, str | None], str | None],
checkpoint_input_dim_reader: Callable[[str], int | None],
entrypoint_log_root: Callable[..., Path],
wrapper_cls: Any = None,
runner_cls: Any = None,
policy_obs_dims_getter: Callable[[Any], tuple[int, int]] | None = None,
sim2sim_preflight: Callable[[str | None], Any] | None = None,
runner_loader: Callable[[Any, str], None] | None = None,
guard_algo_name: str | None = None,
log: LogFn = print,
) -> tuple[RslRlPlaybackSession, str, str | None]:
"""Create a playback session and load the selected policy checkpoint."""
from rsl_rl.runners import OnPolicyRunner
from unilab.rl import RslRlVecEnvAdapter, get_policy_obs_dims
if wrapper_cls is None:
wrapper_cls = RslRlVecEnvAdapter
if runner_cls is None:
runner_cls = OnPolicyRunner
if policy_obs_dims_getter is None:
policy_obs_dims_getter = get_policy_obs_dims
device_name = default_device(torch) if device is None else str(device)
env = env_factory(int(playback_cfg.num_envs))
if env is None:
raise RuntimeError("Playback env factory did not return an environment.")
actor_obs_dim, flat_obs_dim = policy_obs_dims_getter(env.obs_groups_spec)
policy_obs_mode = playback_cfg.policy_obs_mode
checkpoint_path: str | None = None
if playback_cfg.action_mode == "policy":
checkpoint_path = checkpoint_resolver(
playback_cfg.task,
playback_cfg.load_run,
playback_cfg.checkpoint,
playback_cfg.algo_log_name,
playback_cfg.log_root,
)
if policy_obs_mode == "auto" and checkpoint_path is not None:
ckpt_dim = checkpoint_input_dim_reader(checkpoint_path)
if ckpt_dim == actor_obs_dim:
policy_obs_mode = "actor"
elif ckpt_dim == flat_obs_dim:
policy_obs_mode = "flat"
elif ckpt_dim is not None:
raise RuntimeError(
"Checkpoint actor input dim mismatch: "
f"ckpt={ckpt_dim}, actor_obs={actor_obs_dim}, flat_obs={flat_obs_dim}. "
"Please pass --policy_obs_mode actor|flat explicitly if needed."
)
else:
policy_obs_mode = "flat"
wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode)
log(f"Policy obs mode: {policy_obs_mode} (actor_obs={actor_obs_dim}, flat_obs={flat_obs_dim})")
train_cfg = copy.deepcopy(algo_config)
if "runner" not in train_cfg:
train_cfg["runner"] = {}
train_cfg["runner"]["logger"] = "none"
policy = None
runner = None
if playback_cfg.action_mode == "policy":
if checkpoint_path is None:
log("WARNING: no checkpoint found - falling back to zero actions.")
else:
if sim2sim_preflight is not None:
sim2sim_preflight(str(Path(checkpoint_path).parent))
log_dir = str(
entrypoint_log_root(
Path(root_dir),
algo_log_name=playback_cfg.algo_log_name,
log_root=playback_cfg.log_root,
)
/ playback_cfg.task
/ "play_temp"
)
runner = runner_cls(wrapped_env, train_cfg, log_dir=log_dir, device=device_name)
policy_obs_dim = actor_obs_dim if policy_obs_mode == "actor" else flat_obs_dim
action_shape = env.action_space.shape
policy_action_dim = int(action_shape[0]) if action_shape is not None else None
with policy_load_dim_guard(
env_obs_dim=policy_obs_dim,
env_action_dim=policy_action_dim,
algo_name=guard_algo_name or playback_cfg.algo_log_name,
):
if runner_loader is not None:
runner_loader(runner, checkpoint_path)
else:
runner.load(
checkpoint_path,
load_cfg={
"actor": True,
"critic": False,
"optimizer": False,
"iteration": False,
"rnd": False,
},
)
policy = runner.get_inference_policy(device=device_name)
log(f"Action mode: {playback_cfg.action_mode}")
session = RslRlPlaybackSession(
env=env,
wrapped_env=wrapped_env,
device=device_name,
action_mode=playback_cfg.action_mode,
policy=policy,
num_envs=playback_cfg.num_envs,
runner=runner,
)
return session, policy_obs_mode, checkpoint_path
def _cfg_checkpoint_value(cfg: Any) -> str | None:
from omegaconf import OmegaConf
return normalize_checkpoint_value(OmegaConf.select(cfg, "algo.checkpoint", default=None))
def _resolve_appo_checkpoint_from_cfg(
cfg: Any,
*,
root_dir: str | Path,
) -> tuple[str | None, str | None]:
from unilab.training import get_log_root
from unilab.utils.checkpoint import resolve_checkpoint_path
task_log_root = get_log_root(root_dir, cfg) / str(cfg.training.task_name)
selected_path, selected_dir = resolve_checkpoint_path(
task_log_root,
str(cfg.algo.load_run),
checkpoint=_cfg_checkpoint_value(cfg),
)
return (
str(selected_path) if selected_path is not None else None,
str(selected_dir) if selected_dir is not None else None,
)
def _build_appo_actor(
*,
env: Any,
wrapped_env: Any,
cfg: Any,
rl_cfg: dict[str, Any],
device: str,
) -> tuple[Any, int, int]:
"""Build the APPO actor and return it with ``(obs_dim, action_dim)``."""
from copy import deepcopy
from rsl_rl.utils import resolve_callable
from tensordict import TensorDict
from unilab.rl.contracts.observations import get_obs_dims
action_shape = env.action_space.shape
if action_shape is None:
raise ValueError("env.action_space.shape must be defined")
action_dim = int(action_shape[0])
rl_cfg_dict = deepcopy(rl_cfg)
obs_dim, critic_dim = get_obs_dims(env.obs_groups_spec)
num_envs = int(getattr(wrapped_env, "num_envs", getattr(env, "num_envs", 1)))
obs_groups = rl_cfg_dict.setdefault("obs_groups", {})
if "obs_groups" not in rl_cfg_dict or not isinstance(obs_groups, dict):
obs_groups = {}
rl_cfg_dict["obs_groups"] = obs_groups
actor_group = obs_groups.get("actor", obs_groups.get("policy", {}))
if isinstance(actor_group, dict) and "policy" in actor_group:
actor_group["policy"] = obs_dim
obs_groups["actor"] = actor_group
else:
obs_groups["actor"] = {"policy": obs_dim}
critic_group = obs_groups.get("critic")
if critic_group is None:
obs_groups["critic"] = {"policy": critic_dim if critic_dim > 0 else obs_dim}
elif isinstance(critic_group, dict) and "policy" in critic_group:
critic_group["policy"] = critic_dim if critic_dim > 0 else obs_dim
obs_example = torch.zeros((num_envs, obs_dim), device=device)
td_example = TensorDict({"policy": obs_example}, batch_size=num_envs)
actor_cfg = deepcopy(rl_cfg_dict["actor"])
actor_cls = resolve_callable(actor_cfg.pop("class_name"))
actor_cfg.pop("num_actions", None)
actor = actor_cls(td_example, rl_cfg_dict["obs_groups"], "actor", action_dim, **actor_cfg)
return actor.to(device).eval(), obs_dim, action_dim
[docs]
def create_appo_playback_session(
*,
playback_cfg: PlaybackConfig,
cfg: Any,
rl_cfg: dict[str, Any],
env_factory: Callable[[int], Any],
root_dir: str | Path,
device: str | None,
wrapper_cls: Any,
log: LogFn = print,
) -> tuple[RslRlPlaybackSession, str, str | None]:
"""Create an APPO interactive playback session."""
device_name = default_device(torch) if device is None else str(device)
env = env_factory(int(playback_cfg.num_envs))
if env is None:
raise RuntimeError("Playback env factory did not return an environment.")
policy_obs_mode = playback_cfg.policy_obs_mode
wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode)
policy = None
actor = None
checkpoint_path: str | None = None
if playback_cfg.action_mode == "policy":
checkpoint_path, checkpoint_dir = _resolve_appo_checkpoint_from_cfg(cfg, root_dir=root_dir)
if checkpoint_path is None or not Path(checkpoint_path).exists():
log(
"WARNING: no APPO checkpoint found for "
f"load_run={cfg.algo.load_run} - falling back to zero actions."
)
else:
resolve_sim2sim_config(
checkpoint_dir,
cfg,
algo_name="appo",
strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
)
actor, actor_obs_dim, action_dim = _build_appo_actor(
env=env,
wrapped_env=wrapped_env,
cfg=cfg,
rl_cfg=rl_cfg,
device=device_name,
)
checkpoint = torch.load(checkpoint_path, map_location=device_name, weights_only=True)
with policy_load_dim_guard(
env_obs_dim=actor_obs_dim,
env_action_dim=action_dim,
algo_name="appo",
):
actor.load_state_dict(checkpoint["actor"])
policy = actor
log(f"Loading APPO checkpoint: {checkpoint_path}")
log(f"Action mode: {playback_cfg.action_mode}")
return (
RslRlPlaybackSession(
env=env,
wrapped_env=wrapped_env,
device=device_name,
action_mode=playback_cfg.action_mode,
policy=policy,
num_envs=playback_cfg.num_envs,
actor=actor,
),
policy_obs_mode,
checkpoint_path,
)
[docs]
def default_device(torch_module, preferred: str | None = None) -> str:
"""Resolve runtime device with optional user override."""
if preferred:
return preferred
if torch_module.cuda.is_available():
return "cuda"
xpu = getattr(torch_module, "xpu", None)
xpu_is_available = getattr(xpu, "is_available", None)
if callable(xpu_is_available) and xpu_is_available():
return "xpu"
if torch_module.backends.mps.is_available():
return "mps"
return "cpu"
[docs]
def resolve_play_obs_dims(obs_groups_spec: dict[str, int]) -> tuple[int, int]:
from unilab.rl.contracts.observations import get_obs_dims
obs_dim, critic_obs_dim = get_obs_dims(obs_groups_spec)
return int(obs_dim), int(critic_obs_dim)
[docs]
def resolve_playback_device(env_device: Any, requested_device: str) -> torch.device:
"""Resolve actor placement from the authoritative TorchEnv device."""
if env_device is None:
return torch.device(requested_device)
return torch.device(env_device)
[docs]
def resolve_play_actor_spec(
algo_name: str,
cfg: DictConfig,
*,
obs_dim: int,
critic_obs_dim: int,
) -> tuple[str, dict[str, Any]]:
"""Resolve the actor implementation and model kwargs used by off-policy play."""
if algo_name != "sac":
return algo_name, {}
from unilab.rl.runtime.offpolicy.resolution import resolve_custom_offpolicy_runtime
rl_cfg = cast(dict[str, Any], OmegaConf.to_container(cfg.algo, resolve=True))
custom_runtime = resolve_custom_offpolicy_runtime(rl_cfg)
if custom_runtime is None:
return "sac", {}
actor_algo_type = str(custom_runtime.algo_type or algo_name)
actor_kwargs = custom_runtime.build_model_kwargs(
obs_dim=int(obs_dim),
critic_obs_dim=int(critic_obs_dim),
)
return actor_algo_type, actor_kwargs
def _flashsac_actor_kwargs(cfg: DictConfig) -> dict[str, Any]:
kwargs = {
"actor_num_blocks": cfg.algo.algo_params.actor_num_blocks,
"actor_noise_zeta_mu": cfg.algo.algo_params.actor_noise_zeta_mu,
"actor_noise_zeta_max": cfg.algo.algo_params.actor_noise_zeta_max,
}
actor_config = cfg.algo.algo_params.get("actor_config")
if actor_config is not None:
kwargs["actor_config"] = actor_config
return kwargs
[docs]
def build_play_actor(
algo_name: str,
cfg: DictConfig,
*,
obs_dim: int,
critic_obs_dim: int,
action_dim: int,
device: str,
) -> tuple[Any, Any | None, str, dict[str, Any]]:
"""Build the policy actor selected by an off-policy owner config."""
from unilab.rl.algos.common.actor_factory import build_actor
actor_algo_type, actor_kwargs = resolve_play_actor_spec(
algo_name,
cfg,
obs_dim=obs_dim,
critic_obs_dim=critic_obs_dim,
)
normalizer = None
if algo_name == "sac":
actor = build_actor(
actor_algo_type,
obs_dim,
action_dim,
cfg.algo.actor_hidden_dim,
cfg.algo.use_layer_norm,
device,
**actor_kwargs,
)
elif algo_name == "flashsac":
build_actor_kwargs = _flashsac_actor_kwargs(cfg)
actor = build_actor(
algo_name,
obs_dim,
action_dim,
cfg.algo.actor_hidden_dim,
cfg.algo.use_layer_norm,
device,
**build_actor_kwargs,
)
if cfg.algo.obs_normalization:
from unilab.rl.algos.common.normalization import EmpiricalNormalization
normalizer = EmpiricalNormalization(shape=obs_dim, device=device)
else:
raise ValueError(f"Unsupported algo: {algo_name}")
actor.eval()
return actor, normalizer, actor_algo_type, actor_kwargs
[docs]
def flashsac_play_actor_semantics(
cfg: Any,
*,
obs_dim: int,
action_dim: int,
actor: Any | None = None,
) -> dict[str, Any]:
"""Recompute the FlashSAC actor semantics payload expected at play time."""
from unilab.rl.algos.flash_sac.actor_spec import (
FlashSACActorSpec,
build_builtin_actor_semantics,
build_custom_actor_semantics,
)
actor_config = cfg.algo.algo_params.get("actor_config")
if actor_config is not None:
transform_id = getattr(actor, "action_transform_id", None)
if not isinstance(transform_id, str) or not transform_id:
raise TypeError(
"Custom FlashSAC play actor must declare a non-empty string 'action_transform_id'"
)
return build_custom_actor_semantics(
FlashSACActorSpec.from_config(actor_config),
obs_dim=obs_dim,
action_dim=action_dim,
action_transform_id=transform_id,
)
return build_builtin_actor_semantics(
obs_dim=obs_dim,
action_dim=action_dim,
num_blocks=int(cfg.algo.algo_params.actor_num_blocks),
hidden_dim=int(cfg.algo.actor_hidden_dim),
noise_zeta_mu=float(cfg.algo.algo_params.actor_noise_zeta_mu),
noise_zeta_max=int(cfg.algo.algo_params.actor_noise_zeta_max),
)
[docs]
def load_play_actor(
algo_name: str,
actor: Any,
normalizer: Any | None,
checkpoint: dict[str, Any],
expected_actor_semantics: dict[str, Any] | None = None,
) -> None:
"""Restore an off-policy play actor and its optional observation normalizer.
For FlashSAC, ``expected_actor_semantics`` enables fail-closed validation
of the checkpoint's actor construction/observation/action semantics against
the owner config before weights are loaded.
"""
if algo_name == "flashsac":
if expected_actor_semantics is not None:
from unilab.rl.algos.flash_sac.actor_spec import validate_actor_semantics
validate_actor_semantics(checkpoint.get("actor_semantics"), expected_actor_semantics)
actor.load_state_dict(checkpoint["actor"])
elif algo_name == "sac":
actor.load_state_dict(checkpoint["actor"])
else:
raise ValueError(f"Unsupported algo: {algo_name}")
if normalizer is not None and checkpoint.get("obs_normalizer"):
normalizer.load_state_dict(checkpoint["obs_normalizer"])
normalizer.eval()
[docs]
def build_offpolicy_env_cfg_override(
algo_name: str,
cfg: DictConfig,
*,
root_dir: str | Path,
) -> dict[str, Any] | None:
"""Build the task env override for off-policy training."""
from unilab.base.config_adapter import EnvConfigAdapter
from unilab.training.common import assert_offpolicy_task_choice_matches_algo
assert_offpolicy_task_choice_matches_algo(cfg, algo_name=algo_name)
return cast(
dict[str, Any] | None,
EnvConfigAdapter(cfg, root_dir=root_dir, algo_name=algo_name).build_task_env_cfg_override(),
)
[docs]
def build_offpolicy_play_env_cfg_override(
algo_name: str,
cfg: DictConfig,
*,
root_dir: str | Path,
) -> dict[str, Any] | None:
"""Build the off-policy play override, including backend render intent.
Training and playback deliberately use separate adapters: an IsaacSim
renderer must be selected before the worker's Kit ``INIT`` handshake, but
learner/collector environments must stay on the no-rendering experience.
"""
from unilab.base.config_adapter import EnvConfigAdapter
from unilab.training.common import assert_offpolicy_task_choice_matches_algo
assert_offpolicy_task_choice_matches_algo(cfg, algo_name=algo_name)
return cast(
dict[str, Any] | None,
EnvConfigAdapter(cfg, root_dir=root_dir, algo_name=algo_name).build_play_env_cfg_override(),
)
[docs]
def create_sac_playback_session(
*,
playback_cfg: PlaybackConfig,
cfg: Any,
env_factory: Callable[[int], Any],
root_dir: str | Path,
device: str | None,
algo_name: str = "sac",
log: LogFn = print,
) -> tuple[OffPolicyPlaybackSession, str, str | None]:
"""Create an interactive playback session for off-policy actors."""
import os
from unilab.rl.algos.common.actor_factory import build_actor
from unilab.training import get_log_root
from unilab.utils.checkpoint import resolve_checkpoint_path
device_name = default_device(torch, str(device) if device is not None else None)
env = env_factory(int(playback_cfg.num_envs))
if env is None:
raise RuntimeError("Playback env factory did not return an environment.")
obs_dim, critic_obs_dim = resolve_play_obs_dims(env.obs_groups_spec)
action_shape = env.action_space.shape
if action_shape is None:
raise ValueError("env.action_space.shape must be defined")
action_dim = int(action_shape[0])
session_device = resolve_playback_device(getattr(env, "device", None), device_name)
actor_algo_type, actor_kwargs = resolve_play_actor_spec(
algo_name,
cfg,
obs_dim=obs_dim,
critic_obs_dim=critic_obs_dim,
)
if algo_name == "flashsac":
actor_kwargs.update(_flashsac_actor_kwargs(cfg))
actor = None
checkpoint_path: str | None = None
normalizer = None
if bool(getattr(cfg.algo, "obs_normalization", False)):
from unilab.rl.algos.common.normalization import EmpiricalNormalization
normalizer = EmpiricalNormalization(shape=obs_dim, device=session_device)
if playback_cfg.action_mode == "policy":
actor = build_actor(
actor_algo_type,
obs_dim,
action_dim,
cfg.algo.actor_hidden_dim,
cfg.algo.use_layer_norm,
str(session_device),
**actor_kwargs,
)
actor.eval()
algo_log_name = str(cfg.algo.algo_log_name)
if playback_cfg.log_root is not None:
task_log_root = (
get_log_root(
Path(root_dir),
algo_log_name=algo_log_name,
log_root=playback_cfg.log_root,
)
/ playback_cfg.task
)
else:
# Without an explicit log_root the lookup stays on the repo-rooted
# ``<root_dir>/logs/<algo>`` tree; the test-only
# ``UNILAB_TEST_LOG_ROOT`` override is reserved for cfg-driven
# resolution (``training.resolve_configured_checkpoint_path``).
task_log_root = (Path(root_dir) / "logs" / algo_log_name).absolute() / playback_cfg.task
resolved_path, resolved_dir = resolve_checkpoint_path(
task_log_root,
playback_cfg.load_run,
checkpoint=normalize_checkpoint_value(playback_cfg.checkpoint),
)
checkpoint_path = str(resolved_path) if resolved_path is not None else None
checkpoint_dir = str(resolved_dir) if resolved_dir is not None else None
if checkpoint_path is None or not os.path.exists(checkpoint_path):
log(
f"WARNING: no {algo_name} checkpoint found for "
f"load_run={playback_cfg.load_run} - falling back to zero actions."
)
actor = None
else:
resolve_sim2sim_config(
checkpoint_dir,
cfg,
algo_name=algo_name,
strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
)
checkpoint = torch.load(checkpoint_path, map_location=session_device, weights_only=True)
expected_actor_semantics = (
flashsac_play_actor_semantics(
cfg,
obs_dim=obs_dim,
action_dim=action_dim,
actor=actor,
)
if algo_name == "flashsac"
else None
)
with policy_load_dim_guard(
env_obs_dim=obs_dim,
env_action_dim=action_dim,
algo_name=algo_name,
):
load_play_actor(
algo_name,
actor,
normalizer,
checkpoint,
expected_actor_semantics=expected_actor_semantics,
)
log(f"Loading {algo_name} checkpoint: {checkpoint_path}")
log(f"Action mode: {playback_cfg.action_mode}")
return (
OffPolicyPlaybackSession(
env=env,
device=device_name,
action_mode=playback_cfg.action_mode,
actor=actor,
actor_algo_type=actor_algo_type,
normalizer=normalizer,
num_envs=playback_cfg.num_envs,
obs_extractor=extract_play_obs,
),
"actor",
checkpoint_path,
)
[docs]
def prepare_motion_overlay_selection(
env: Any,
*,
show_target_bodies: bool,
show_reward_debug: bool,
target_body_names: str,
target_max_bodies: int,
log: LogFn = print,
) -> MotionOverlaySelection:
"""Resolve body indices used by motion-target and reward-debug overlays."""
if not (show_target_bodies or show_reward_debug):
return MotionOverlaySelection(
enabled=False,
selected_indices=np.zeros((0,), dtype=np.int32),
)
if not (hasattr(env, "motion_loader") and hasattr(env, "motion_sampler")):
log("WARNING: target/reward visualization only works for motion-tracking tasks.")
return MotionOverlaySelection(
enabled=False,
selected_indices=np.zeros((0,), dtype=np.int32),
)
names = tuple(getattr(env.cfg, "body_names", ()))
if len(names) == 0:
log("WARNING: task has no body_names; cannot visualize targets.")
return MotionOverlaySelection(
enabled=False,
selected_indices=np.zeros((0,), dtype=np.int32),
)
name_to_idx = {name: i for i, name in enumerate(names)}
if target_body_names.strip():
chosen = []
for name in [n.strip() for n in target_body_names.split(",") if n.strip()]:
if name in name_to_idx:
chosen.append(name_to_idx[name])
else:
log(f"WARNING: body name not found in task body list: {name}")
selected_indices = np.array(chosen, dtype=np.int32)
else:
selected_indices = np.arange(len(names), dtype=np.int32)
if target_max_bodies > 0:
selected_indices = selected_indices[:target_max_bodies]
return MotionOverlaySelection(
enabled=selected_indices.size > 0,
selected_indices=selected_indices,
)
__all__ = [
"KeyboardCommander",
"MotionOverlaySelection",
"OffPolicyPlaybackSession",
"PlayInteractiveArgs",
"PlaybackControls",
"PlaybackSession",
"PlaybackConfig",
"RslRlPlaybackSession",
"available_backends_for_task",
"build_play_backend_adapter",
"build_playback_config",
"create_appo_playback_session",
"create_rsl_rl_playback_session",
"create_sac_playback_session",
"default_device",
"flashsac_play_actor_semantics",
"infer_checkpoint_actor_input_dim",
"make_sim2sim_preflight",
"normalize_checkpoint_value",
"prepare_motion_overlay_selection",
]