diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index e552ec2cc..cd7510647 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -61,16 +61,22 @@ get_policy_obs_dims, normalize_ppo_train_cfg, ) +from unilab.utils.rotation import np_matrix_from_quat from unilab.visualization.interactive_playback import ( _HORA_DISTILL_CHECKPOINT_UNAVAILABLE, KeyboardCommander, PlaybackControls, - RslRlPlaybackConfig, + PlayInteractiveArgs, + available_backends_for_task, + build_play_backend_adapter, + build_playback_config, create_appo_playback_session, create_hora_distill_playback_session, create_rsl_rl_playback_session, create_sac_playback_session, + infer_checkpoint_actor_input_dim, make_sim2sim_preflight, + normalize_checkpoint_value, prepare_motion_overlay_selection, select_torch_device, ) @@ -108,74 +114,6 @@ import mujoco.viewer -@dataclass -class PlayInteractiveArgs: - 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" - - -def _infer_checkpoint_actor_input_dim(ckpt_path: str) -> int | None: - loaded = torch.load(ckpt_path, map_location="cpu", weights_only=True) - state_dict = loaded.get("actor_state_dict") - if not isinstance(state_dict, dict): - return None - - # Common rsl-rl naming: "mlp.0.weight" or nested prefixes ending with ".0.weight". - for key in ("mlp.0.weight", "actor.mlp.0.weight"): - w = state_dict.get(key) - if isinstance(w, torch.Tensor) and w.ndim == 2: - return int(w.shape[1]) - - for key, w in state_dict.items(): - if key.endswith(".0.weight") and isinstance(w, torch.Tensor) and w.ndim == 2: - return int(w.shape[1]) - return None - - -def _backend_adapter(cfg: DictConfig, *, algo_name: str = "ppo"): - from unilab.base.backend import materialize_scene_visual_override - from unilab.training import BackendAdapter - - return BackendAdapter( - cfg, - root_dir=ROOT_DIR, - algo_name=algo_name, - scene_materializer=materialize_scene_visual_override, - ) - - def _algo_config_dict(cfg: DictConfig | None) -> dict[str, Any]: """Return the composed PPO algo config as a plain dict. @@ -331,15 +269,7 @@ def _quat_to_rotmat_wxyz(quat: np.ndarray) -> np.ndarray: n = np.linalg.norm(q) if n < 1e-12: return np.eye(3, dtype=np.float64) - w, x, y, z = q / n - return np.array( - [ - [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - z * w), 2.0 * (x * z + y * w)], - [2.0 * (x * y + z * w), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - x * w)], - [2.0 * (x * z - y * w), 2.0 * (y * z + x * w), 1.0 - 2.0 * (x * x + y * y)], - ], - dtype=np.float64, - ) + return np_matrix_from_quat(q / n) def _add_sphere_marker(scene, pos: np.ndarray, radius: float, rgba: np.ndarray) -> bool: @@ -462,15 +392,6 @@ def _default_viewer_camera_distance(mj_model, env: Any, *, follow_body: bool) -> return min(extent_distance, _FOLLOW_CAMERA_MAX_DISTANCE) -def _available_backends_for_task(task_name: str) -> tuple[str, ...]: - 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) - - def _can_launch_glfw_viewer() -> bool: try: import glfw @@ -823,21 +744,6 @@ def _load_viewer_model(env: Any, *, use_env_visual_model: bool): return playback_model -def _build_playback_config(args, *, num_envs: int = 1) -> RslRlPlaybackConfig: - return RslRlPlaybackConfig( - 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)), - ) - - def _build_keyboard_commander(env: Any, args) -> KeyboardCommander | None: """Set up keyboard velocity teleop, or return None when unsupported/disabled.""" if not bool(getattr(args, "keyboard", False)): @@ -989,7 +895,7 @@ def play_interactive(args, cfg: DictConfig | None = None, *, algo: str | None = algo = str(algo or getattr(args, "algo", "ppo")) # Always use a single env for interactive view - available_backends = _available_backends_for_task(args.task) + available_backends = available_backends_for_task(args.task) if available_backends and "mujoco" not in available_backends: print( "[play_interactive] Task does not support MuJoCo backend: " @@ -1006,7 +912,9 @@ def _create_env(num_envs: int): if algo in _OFFPOLICY_INTERACTIVE_ALGOS: env_cfg_override = build_offpolicy_env_cfg_override(algo, cfg, root_dir=ROOT_DIR) else: - env_cfg_override = _backend_adapter(cfg, algo_name=algo).build_task_env_cfg_override() + env_cfg_override = build_play_backend_adapter( + cfg, root_dir=ROOT_DIR, algo_name=algo + ).build_task_env_cfg_override() try: return create_env( cfg, @@ -1026,7 +934,7 @@ def _create_env(num_envs: int): raise try: - playback_cfg = _build_playback_config(args, num_envs=1) + playback_cfg = build_playback_config(args, num_envs=1) if algo == "ppo": wrapper_cls = RslRlVecEnvWrapper if cfg is not None: @@ -1043,7 +951,7 @@ def _create_env(num_envs: int): root_dir=ROOT_DIR, device=device, checkpoint_resolver=resolve_checkpoint, - checkpoint_input_dim_reader=_infer_checkpoint_actor_input_dim, + checkpoint_input_dim_reader=infer_checkpoint_actor_input_dim, entrypoint_log_root=get_entrypoint_log_root, wrapper_cls=wrapper_cls, runner_cls=OnPolicyRunner, @@ -1288,18 +1196,11 @@ def _on_key(keycode: int) -> None: print("[play_interactive] Done.") -def _normalize_checkpoint_value(value: object) -> str | None: - if value is None: - return None - text = str(value) - return None if text in {"-1", "None", "null"} else text - - def _build_play_args(cfg: DictConfig, *, algo: str = "ppo") -> PlayInteractiveArgs: return PlayInteractiveArgs( task=str(cfg.training.task_name), load_run=str(cfg.algo.load_run), - checkpoint=_normalize_checkpoint_value(OmegaConf.select(cfg, "algo.checkpoint")), + checkpoint=normalize_checkpoint_value(OmegaConf.select(cfg, "algo.checkpoint")), action_mode=str(cfg.interactive.action_mode), policy_obs_mode=str(cfg.interactive.policy_obs_mode), algo_log_name=str(cfg.algo.algo_log_name), diff --git a/scripts/play_viser.py b/scripts/play_viser.py index 911abe4af..4f6405477 100644 --- a/scripts/play_viser.py +++ b/scripts/play_viser.py @@ -57,7 +57,12 @@ ) from unilab.visualization.interactive_playback import ( PlaybackControls, + PlayInteractiveArgs, + available_backends_for_task, + build_play_backend_adapter, + build_playback_config, create_rsl_rl_playback_session, + infer_checkpoint_actor_input_dim, make_sim2sim_preflight, select_torch_device, ) @@ -84,14 +89,7 @@ import mujoco import viser # noqa: E402 -from play_interactive import ( # noqa: E402 - PlayInteractiveArgs, - _available_backends_for_task, - _backend_adapter, - _build_playback_config, - _infer_checkpoint_actor_input_dim, - resolve_checkpoint, -) +from play_interactive import resolve_checkpoint # noqa: E402 from unilab.training import algo_config_dict # noqa: E402 @@ -206,7 +204,7 @@ def play_viser(args: PlayInteractiveArgs, cfg: DictConfig) -> None: print(f"[play_viser] Device: {device}") # --- Validate backend --------------------------------------------------- - available_backends = _available_backends_for_task(args.task) + available_backends = available_backends_for_task(args.task) if available_backends and "mujoco" not in available_backends: print( f"[play_viser] Task {args.task} does not support MuJoCo backend. " @@ -222,7 +220,9 @@ def _create_env(env_count: int): return registry.make(args.task, num_envs=env_count, sim_backend="mujoco") from unilab.training import create_env - env_cfg_override = _backend_adapter(cfg).build_task_env_cfg_override() + env_cfg_override = build_play_backend_adapter( + cfg, root_dir=ROOT_DIR + ).build_task_env_cfg_override() return create_env( cfg, num_envs=env_count, @@ -232,13 +232,13 @@ def _create_env(env_count: int): ) playback_session, _policy_obs_mode, _checkpoint_path = create_rsl_rl_playback_session( - playback_cfg=_build_playback_config(args, num_envs=num_envs), + playback_cfg=build_playback_config(args, num_envs=num_envs), env_factory=_create_env, algo_config=algo_config_dict(cfg), root_dir=ROOT_DIR, device=device, checkpoint_resolver=resolve_checkpoint, - checkpoint_input_dim_reader=_infer_checkpoint_actor_input_dim, + checkpoint_input_dim_reader=infer_checkpoint_actor_input_dim, entrypoint_log_root=get_entrypoint_log_root, wrapper_cls=RslRlVecEnvWrapper, runner_cls=OnPolicyRunner, diff --git a/scripts/train_appo.py b/scripts/train_appo.py index 9f94f55b5..bbd6aa1d7 100644 --- a/scripts/train_appo.py +++ b/scripts/train_appo.py @@ -7,7 +7,7 @@ import sys from collections.abc import Callable from pathlib import Path -from typing import Any, cast +from typing import Any import hydra import torch @@ -32,7 +32,12 @@ ) from unilab.training.experiment import ExperimentTracker from unilab.training.onnx_export import export_policy_onnx, verify_policy_onnx -from unilab.training.sim2sim import policy_load_dim_guard, resolve_sim2sim_config +from unilab.training.rsl_rl import RslRlVecEnvWrapper +from unilab.visualization.interactive_playback import ( + RslRlPlaybackConfig, + create_appo_playback_session, + normalize_checkpoint_value, +) def _training_resume_requested(load_run: Any) -> bool: @@ -155,9 +160,6 @@ def play_appo( native Motrix viewer or when no checkpoint could be resolved. """ del root_dir - import numpy as np - from rsl_rl.utils import resolve_callable - from tensordict import TensorDict if resolve_checkpoint_path is not None: load_path, load_path_dir = resolve_checkpoint_path(cfg) @@ -170,20 +172,6 @@ def play_appo( print(f"Could not find run to load. load_path={load_path}") return None - cfg = ( - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name="appo", - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - or cfg - ) - - env_cfg_override = BackendAdapter( - cfg, root_dir=ROOT_DIR, algo_name="appo" - ).build_task_env_cfg_override() - device = cfg.training.device or ( "cuda" if torch.cuda.is_available() @@ -193,58 +181,37 @@ def play_appo( ) print(f"Using device for play: {device}") - env = cast( - Any, - create_env( + playback_cfg = RslRlPlaybackConfig( + task=str(cfg.training.task_name), + load_run=str(cfg.algo.load_run), + checkpoint=normalize_checkpoint_value( + OmegaConf.select(cfg, "algo.checkpoint", default=None) + ), + action_mode="policy", + policy_obs_mode="flat", + algo_log_name=str(cfg.algo.algo_log_name), + log_root=None, + num_envs=cfg.training.play_env_num, + ) + session, _policy_obs_mode, _checkpoint_path = create_appo_playback_session( + playback_cfg=playback_cfg, + cfg=cfg, + rl_cfg=rl_cfg, + env_factory=lambda n: create_env( cfg, - num_envs=cfg.training.play_env_num, - env_cfg_override=env_cfg_override, + num_envs=n, + env_cfg_override=BackendAdapter( + cfg, root_dir=ROOT_DIR, algo_name="appo" + ).build_task_env_cfg_override(), ), + root_dir=ROOT_DIR, + device=device, + wrapper_cls=RslRlVecEnvWrapper, ) - from unilab.base.observations import get_obs_dims - - obs_dim, critic_dim = get_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]) - - rl_cfg_dict = dict(rl_cfg) - if "obs_groups" not in rl_cfg_dict: - rl_cfg_dict["obs_groups"] = { - "actor": {"policy": obs_dim}, - "critic": {"policy": critic_dim if critic_dim > 0 else obs_dim}, - } - else: - actor_group = rl_cfg_dict["obs_groups"].get( - "actor", rl_cfg_dict["obs_groups"].get("policy", {}) - ) - if isinstance(actor_group, dict) and "policy" in actor_group: - actor_group["policy"] = obs_dim - critic_group = rl_cfg_dict["obs_groups"].get("critic") - if critic_group is None: - rl_cfg_dict["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 - - from copy import deepcopy - - obs_example = torch.zeros((cfg.training.play_env_num, obs_dim), device=device) - td_example = TensorDict({"policy": obs_example}, batch_size=cfg.training.play_env_num) - - 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) - actor = actor.to(device) - actor.eval() - - print(f"Loading model: {load_path}") - checkpoint = torch.load(load_path, map_location=device, weights_only=True) - with policy_load_dim_guard(env_obs_dim=obs_dim, env_action_dim=action_dim, algo_name="appo"): - actor.load_state_dict(checkpoint["actor"]) + env = session.env + actor = session.actor + # The checkpoint early-return above guarantees a loaded actor here. + assert actor is not None # Export actor to ONNX if load_path_dir is not None: @@ -260,6 +227,7 @@ def forward(self, obs: torch.Tensor) -> torch.Tensor: export_module = _DeterministicAPPOActor(actor.mlp) onnx_path = os.path.join(load_path_dir, "policy.onnx") + obs_dim = int(session.wrapped_env.num_obs) dummy_input = torch.randn(1, obs_dim, device=device) export_policy_onnx(export_module, onnx_path, (dummy_input,), input_names=["obs"]) @@ -267,9 +235,6 @@ def forward(self, obs: torch.Tensor) -> torch.Tensor: verify_input = torch.randn(1, obs_dim, device=device) verify_policy_onnx(export_module, onnx_path, (verify_input,), input_names=["obs"]) - if env.state is None: - env.init_state() - with torch.inference_mode(): play_video_path = env.run_playback_mode( play_render_mode=getattr(cfg.training, "play_render_mode", "auto"), @@ -278,24 +243,8 @@ def forward(self, obs: torch.Tensor) -> torch.Tensor: render_spacing=float( getattr(cfg.training, "render_spacing", getattr(env.cfg, "render_spacing", 1.0)) ), - initialize=lambda: np.asarray( - env.reset(np.arange(cfg.training.play_env_num, dtype=np.int32))[0]["obs"], - dtype=np.float32, - ), - step=lambda obs_np: np.asarray( - env.step( - actor( - TensorDict( - {"policy": torch.from_numpy(obs_np).to(device)}, - batch_size=cfg.training.play_env_num, - ) - ) - .cpu() - .numpy() - .astype(np.float32) - ).obs["obs"], - dtype=np.float32, - ), + initialize=session.reset, + step=lambda _obs: session.step_once(), camera_kwargs={ "cam_distance": cfg.training.cam_distance, "cam_elevation": cfg.training.cam_elevation, diff --git a/scripts/train_him_ppo.py b/scripts/train_him_ppo.py index d2b3680e4..e74781fc7 100644 --- a/scripts/train_him_ppo.py +++ b/scripts/train_him_ppo.py @@ -28,12 +28,20 @@ create_env, ensure_registries, format_play_checkpoint_error, + get_entrypoint_log_root, get_log_root, parse_checkpoint_path, ) from unilab.training.experiment import ExperimentTracker -from unilab.training.sim2sim import policy_load_dim_guard, resolve_sim2sim_config +from unilab.training.rsl_rl import RslRlVecEnvWrapper, get_policy_obs_dims from unilab.visualization import render_play_mode +from unilab.visualization.interactive_playback import ( + RslRlPlaybackConfig, + create_rsl_rl_playback_session, + infer_checkpoint_actor_input_dim, + make_sim2sim_preflight, + normalize_checkpoint_value, +) def _backend_adapter(cfg: DictConfig) -> BackendAdapter: @@ -75,31 +83,49 @@ def play_him_ppo(cfg: DictConfig, device: str) -> str | None: ) return None - cfg = ( - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name="ppo", - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - or cfg + def _create_env(num_envs: int): + env_cfg_override = cast(dict[str, Any], _backend_adapter(cfg).build_play_env_cfg_override()) + return create_env(cfg, num_envs=num_envs, env_cfg_override=env_cfg_override) + + session, _policy_obs_mode, _checkpoint_path = create_rsl_rl_playback_session( + playback_cfg=RslRlPlaybackConfig( + task=str(cfg.training.task_name), + load_run=str(getattr(cfg.algo, "load_run", "-1")), + checkpoint=normalize_checkpoint_value(getattr(cfg.algo, "checkpoint", None)), + action_mode="policy", + policy_obs_mode="flat", + algo_log_name=str(cfg.algo.algo_log_name), + log_root=getattr(cfg.training, "log_root", None), + num_envs=int(cfg.training.play_env_num), + ), + env_factory=_create_env, + algo_config=rl_cfg, + root_dir=ROOT_DIR, + device=device, + # The checkpoint was already resolved above for the friendly early exit. + checkpoint_resolver=lambda *_args: str(load_path), + checkpoint_input_dim_reader=infer_checkpoint_actor_input_dim, + entrypoint_log_root=get_entrypoint_log_root, + wrapper_cls=RslRlVecEnvWrapper, + runner_cls=HIMOnPolicyRunner, + # HIMOnPolicyRunner.load does not accept a load_cfg argument. + runner_loader=lambda runner, path: runner.load(path), + policy_obs_dims_getter=get_policy_obs_dims, + train_cfg_normalizer=lambda train_cfg: train_cfg, + sim2sim_preflight=make_sim2sim_preflight(cfg, algo_name="ppo"), + guard_algo_name="him_ppo", ) - env_cfg_override = cast(dict[str, Any], _backend_adapter(cfg).build_play_env_cfg_override()) - env = create_env(cfg, num_envs=cfg.training.play_env_num, env_cfg_override=env_cfg_override) - from unilab.training.rsl_rl import RslRlVecEnvWrapper - - wrapped_env = RslRlVecEnvWrapper(env, device=device) - runner = HIMOnPolicyRunner(wrapped_env, rl_cfg, log_dir=None, device=device) - with policy_load_dim_guard( - env_obs_dim=getattr(wrapped_env, "num_obs", None), - env_action_dim=getattr(wrapped_env, "num_actions", None), - algo_name="him_ppo", - ): - runner.load(str(load_path)) - policy = runner.get_inference_policy(device=device) + env = session.env + assert session.runner is not None and session.policy is not None + + # HIM's inference policy consumes the flat actor tensor, not the full obs + # TensorDict the session hands to ``policy``. + him_policy = session.policy + session.policy = lambda obs: him_policy(obs["actor"]) + if EXPORT_POLICY: - runner.export_policy_to_onnx(path=str(load_path_dir)) - runner.export_policy_to_jit(path=str(load_path_dir)) + session.runner.export_policy_to_onnx(path=str(load_path_dir)) + session.runner.export_policy_to_jit(path=str(load_path_dir)) output_video = Path(load_path_dir) / "play_video.mp4" print(f"Rendering video to {output_video}...") @@ -113,8 +139,8 @@ def play_him_ppo(cfg: DictConfig, device: str) -> str | None: ), num_steps=cfg.training.play_steps, output_video=output_video, - initialize=lambda: wrapped_env.reset()[0]["actor"], - step=lambda obs: wrapped_env.step(policy(obs))[0]["actor"], + initialize=lambda: session.reset()["actor"], + step=lambda _obs: session.step_once()["actor"], camera_kwargs={ "cam_distance": cfg.training.cam_distance, "cam_elevation": cfg.training.cam_elevation, @@ -186,7 +212,6 @@ def main(cfg: DictConfig) -> None: try: if not cfg.training.play_only: env = create_env(cfg, num_envs=cfg.algo.num_envs, env_cfg_override=env_cfg_override) - from unilab.training.rsl_rl import RslRlVecEnvWrapper apply_env_nan_guard(env, cfg.training) diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 5b32993f0..5b7c020db 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -46,19 +46,18 @@ build_offpolicy_env_cfg_override as _build_offpolicy_env_cfg_override, ) from unilab.training.offpolicy import ( - build_play_actor, default_device, - extract_play_obs, - extract_reset_obs, - load_play_actor, - resolve_play_obs_dim, + resolve_play_actor_spec, resolve_play_obs_dims, ) from unilab.training.onnx_export import export_policy_onnx, verify_policy_onnx from unilab.training.run import ( resolve_offpolicy_checkpoint_path as resolve_checkpoint_path, ) -from unilab.training.sim2sim import policy_load_dim_guard, resolve_sim2sim_config +from unilab.visualization.interactive_playback import ( + RslRlPlaybackConfig, + create_sac_playback_session, +) def enable_faulthandler() -> None: @@ -209,11 +208,8 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: """Play pipeline for off-policy algorithms.""" - import numpy as np import torch - from unilab.algos.offpolicy.worker import resolve_offpolicy_actor_priv_info - load_path, load_path_dir = resolve_checkpoint_path( ROOT_DIR, cfg.algo.algo_log_name, @@ -224,55 +220,46 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: print(f"Could not find checkpoint. load_path={load_path}") return None - cfg = ( - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name=algo_name, - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - or cfg - ) - - env_cfg_override = build_offpolicy_env_cfg_override(algo_name, cfg) - devices = resolve_dp_topology(cfg.training.devices) device = default_device(torch, resolve_dp_rank_device(devices, current_dp_rank())) print(f"Using device for play: {device}") - env = cast( - Any, - create_env( + playback_cfg = RslRlPlaybackConfig( + task=str(cfg.training.task_name), + load_run=str(cfg.algo.load_run), + checkpoint=None, + action_mode="policy", + policy_obs_mode="actor", + algo_log_name=str(cfg.algo.algo_log_name), + log_root=None, + num_envs=int(cfg.training.play_env_num), + ) + session, _policy_obs_mode, _checkpoint_path = create_sac_playback_session( + playback_cfg=playback_cfg, + cfg=cfg, + env_factory=lambda n: create_env( cfg, - num_envs=cfg.training.play_env_num, - env_cfg_override=env_cfg_override, + num_envs=n, + env_cfg_override=build_offpolicy_env_cfg_override(algo_name, cfg), ), - ) - 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]) - actor, normalizer, actor_algo_type, actor_kwargs = build_play_actor( - algo_name, - cfg, - obs_dim=obs_dim, - critic_obs_dim=critic_obs_dim, - action_dim=action_dim, + root_dir=ROOT_DIR, device=device, + algo_name=algo_name, ) - print(f"Loading model: {load_path}") - checkpoint = torch.load(load_path, map_location=device, weights_only=True) - 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, - ) + env = cast(Any, session.env) + actor = session.actor + normalizer = session.normalizer + actor_algo_type = session.actor_algo_type # Export actor to ONNX if load_path_dir is not None and bool(getattr(cfg.training, "export_onnx", True)): + obs_dim, critic_obs_dim = resolve_play_obs_dims(env.obs_groups_spec) + _, actor_kwargs = resolve_play_actor_spec( + algo_name, + cfg, + obs_dim=obs_dim, + critic_obs_dim=critic_obs_dim, + ) onnx_path = os.path.join(load_path_dir, "policy.onnx") dummy_input = torch.randn(1, obs_dim, device=device) dummy_priv_info = ( @@ -313,72 +300,13 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: elif load_path_dir is not None: print("Skipping ONNX export because training.export_onnx=false.") - if env.state is None: - env.init_state() - - current_priv_info: np.ndarray | None = None - - def _resolve_play_priv_info(obs_dict: dict[str, np.ndarray], info: dict | None) -> np.ndarray: - if actor_algo_type != "hora_sac": - raise ValueError("Privileged play info was requested for a non-HORA actor.") - from unilab.base.observations import split_obs_dict - - actor_obs_np, critic_np = split_obs_dict(obs_dict) - priv_info = resolve_offpolicy_actor_priv_info( - algo_type=actor_algo_type, - obs_np=np.asarray(actor_obs_np, dtype=np.float32), - critic_np=np.asarray(critic_np, dtype=np.float32), - info=info, - ) - if priv_info is None: - raise ValueError("HORA-SAC play step is missing privileged info.") - return priv_info - - def _extract_reset_play_obs(reset_result) -> np.ndarray: - nonlocal current_priv_info - 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, info_out = reset_result - if actor_algo_type == "hora_sac": - current_priv_info = _resolve_play_priv_info(obs_out, info_out) - return np.asarray(extract_play_obs(obs_out), dtype=np.float32) - - def _policy_step(obs_np: np.ndarray) -> np.ndarray: - nonlocal current_priv_info - obs_torch = torch.from_numpy(obs_np).to(device) - if normalizer: - obs_torch = normalizer(obs_torch, update=False) - if actor_algo_type == "hora_sac": - if current_priv_info is None: - raise ValueError("HORA-SAC play step is missing privileged info.") - priv_info_torch = torch.from_numpy(current_priv_info).to(device) - actions_np = ( - actor.explore( - obs_torch, - priv_info_torch, - deterministic=True, - ) - .cpu() - .numpy() - ) - elif algo_name in ("sac", "flashsac"): - actions_np = actor.explore(obs_torch, deterministic=True).cpu().numpy() - else: - actions_np = actor(obs_torch).cpu().numpy() - state = env.step(actions_np) - if actor_algo_type == "hora_sac": - current_priv_info = _resolve_play_priv_info(state.obs, state.info) - return np.asarray(extract_play_obs(state.obs), dtype=np.float32) - with torch.inference_mode(): play_video_path = env.run_playback_mode( play_render_mode=getattr(cfg.training, "play_render_mode", "auto"), play_steps=getattr(cfg.training, "play_steps", None), output_video=os.path.join(load_path_dir, "play_video.mp4") if load_path_dir else None, - initialize=lambda: _extract_reset_play_obs( - env.reset(np.arange(cfg.training.play_env_num, dtype=np.int32)) - ), - step=_policy_step, + initialize=session.reset, + step=lambda _obs: session.step_once(), camera_kwargs={ "cam_distance": cfg.training.cam_distance, "cam_elevation": cfg.training.cam_elevation, diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index ccab57ac0..20cc8ad87 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -38,6 +38,7 @@ create_env, ensure_registries, format_play_checkpoint_error, + get_entrypoint_log_root, get_log_root, log_playback_plan, parse_checkpoint_path, @@ -53,13 +54,20 @@ RslRlVecEnvWrapper, apply_rsl_rl_rank_seed, finish_rsl_rl_distributed, + get_policy_obs_dims, normalize_ppo_train_cfg, ppo_samples_per_iteration, resolve_rsl_rl_device, rsl_rl_single_process_topology, ) -from unilab.training.sim2sim import policy_load_dim_guard, resolve_sim2sim_config from unilab.utils.device import get_default_device +from unilab.visualization.interactive_playback import ( + RslRlPlaybackConfig, + create_rsl_rl_playback_session, + infer_checkpoint_actor_input_dim, + make_sim2sim_preflight, + normalize_checkpoint_value, +) try: from rsl_rl.runners import OnPolicyRunner @@ -192,7 +200,6 @@ def _resolve_play_num_steps(cfg: DictConfig) -> int | None: def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: """Play mode for RSL-RL.""" rl_cfg = algo_config_dict(cfg) - wrapper_cls = _resolve_ppo_wrapper_cls(rl_cfg) task_log_root = get_log_root(ROOT_DIR, cfg) / str(cfg.training.task_name) load_path, load_path_dir = parse_checkpoint_path(cfg, root_dir=ROOT_DIR) @@ -216,42 +223,48 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: ) return None - cfg = ( - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name="ppo", - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - or cfg - ) - - env_cfg_override = build_ppo_play_env_cfg_override(cfg) - - env = create_env( - cfg, + def _normalize_play_train_cfg(train_cfg: dict[str, Any]) -> dict[str, Any]: + normalized = normalize_ppo_train_cfg(train_cfg) + apply_ppo_runtime_flags(normalized, cfg, training_enabled=False) + return normalized + + playback_cfg = RslRlPlaybackConfig( + task=str(cfg.training.task_name), + load_run=str(cfg.algo.load_run), + checkpoint=normalize_checkpoint_value( + OmegaConf.select(cfg, "algo.checkpoint", default=None) + ), + action_mode="policy", + policy_obs_mode="flat", + algo_log_name=str(cfg.algo.algo_log_name), + log_root=None, num_envs=cfg.training.play_env_num, - env_cfg_override=env_cfg_override, ) - wrapped_env = wrapper_cls(env, device=device) - train_cfg = normalize_ppo_train_cfg(rl_cfg) - apply_ppo_runtime_flags(train_cfg, cfg, training_enabled=False) - if "runner" not in train_cfg: - train_cfg["runner"] = {} - train_cfg["runner"]["logger"] = "none" - - runner = cast( - Any, - OnPolicyRunner(cast(Any, wrapped_env), train_cfg, log_dir=None, device=device), + session, _policy_obs_mode, _checkpoint_path = create_rsl_rl_playback_session( + playback_cfg=playback_cfg, + env_factory=lambda n: create_env( + cfg, + num_envs=n, + env_cfg_override=build_ppo_play_env_cfg_override(cfg), + ), + algo_config=rl_cfg, + root_dir=ROOT_DIR, + device=device, + checkpoint_resolver=lambda *_args: str(load_path), + checkpoint_input_dim_reader=infer_checkpoint_actor_input_dim, + entrypoint_log_root=get_entrypoint_log_root, + wrapper_cls=_resolve_ppo_wrapper_cls(rl_cfg), + runner_cls=OnPolicyRunner, + policy_obs_dims_getter=get_policy_obs_dims, + train_cfg_normalizer=_normalize_play_train_cfg, + sim2sim_preflight=make_sim2sim_preflight(cfg, algo_name="ppo"), + guard_algo_name="ppo", ) - with policy_load_dim_guard( - env_obs_dim=getattr(wrapped_env, "num_obs", None), - env_action_dim=getattr(wrapped_env, "num_actions", None), - algo_name="ppo", - ): - runner.load(str(load_path), map_location=device) - policy = runner.get_inference_policy(device=device) + env = session.env + runner = session.runner if EXPORT_POLICY: + # The checkpoint early-returns above guarantee a loaded runner here. + assert runner is not None runner.export_policy_to_onnx(path=str(load_path_dir)) runner.export_policy_to_jit(path=str(load_path_dir)) num_steps = _resolve_play_num_steps(cfg) @@ -273,8 +286,8 @@ def _log_plan(plan) -> None: getattr(cfg.training, "render_spacing", getattr(env.cfg, "render_spacing", 1.0)) ), render_offset_mode=str(getattr(env.cfg, "render_offset_mode", "grid")), - initialize=lambda: wrapped_env.reset()[0], - step=lambda obs: wrapped_env.step(policy(obs))[0], + initialize=session.reset, + step=lambda _obs: session.step_once(), camera_kwargs={ "cam_distance": cfg.training.cam_distance, "cam_elevation": cfg.training.cam_elevation, diff --git a/src/unilab/visualization/interactive_playback.py b/src/unilab/visualization/interactive_playback.py index b76b8d7c0..9d8722c31 100644 --- a/src/unilab/visualization/interactive_playback.py +++ b/src/unilab/visualization/interactive_playback.py @@ -32,6 +32,62 @@ class RslRlPlaybackConfig: start_paused: bool = False +@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" + + +def build_playback_config(args: Any, *, num_envs: int = 1) -> RslRlPlaybackConfig: + """Build an :class:`RslRlPlaybackConfig` from a play-args-like object.""" + return RslRlPlaybackConfig( + 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)), + ) + + @dataclass class PlaybackControls: """Viewer-independent playback control state.""" @@ -152,6 +208,8 @@ def __init__( 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 @@ -159,6 +217,8 @@ def __init__( 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 @@ -358,6 +418,55 @@ def select_torch_device() -> str: return "cpu" +def infer_checkpoint_actor_input_dim(ckpt_path: str) -> int | None: + """Infer the actor MLP input dim from an rsl_rl checkpoint, if detectable.""" + loaded = torch.load(ckpt_path, map_location="cpu", weights_only=True) + state_dict = loaded.get("actor_state_dict") + if not isinstance(state_dict, dict): + return None + + # Common rsl-rl naming: "mlp.0.weight" or nested prefixes ending with ".0.weight". + for key in ("mlp.0.weight", "actor.mlp.0.weight"): + w = state_dict.get(key) + if isinstance(w, torch.Tensor) and w.ndim == 2: + return int(w.shape[1]) + + for key, w in state_dict.items(): + if key.endswith(".0.weight") and isinstance(w, torch.Tensor) and w.ndim == 2: + return int(w.shape[1]) + return None + + +def _scene_visual_materializer() -> Any: + from unilab.base.backend import materialize_scene_visual_override + + return materialize_scene_visual_override + + +def build_play_backend_adapter(cfg: Any, *, root_dir: str | Path, algo_name: str = "ppo") -> Any: + """Build the BackendAdapter used by play entrypoints to derive env cfg overrides.""" + from unilab.training import BackendAdapter + + return BackendAdapter( + cfg, + root_dir=root_dir, + algo_name=algo_name, + scene_materializer=_scene_visual_materializer(), + ) + + +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) + + def create_rsl_rl_playback_session( *, playback_cfg: RslRlPlaybackConfig, @@ -373,6 +482,8 @@ def create_rsl_rl_playback_session( policy_obs_dims_getter: Callable[[Any], tuple[int, int]], train_cfg_normalizer: Callable[[dict[str, Any]], dict[str, Any]], 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.""" @@ -417,6 +528,7 @@ def create_rsl_rl_playback_session( 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.") @@ -439,18 +551,21 @@ def create_rsl_rl_playback_session( with policy_load_dim_guard( env_obs_dim=policy_obs_dim, env_action_dim=policy_action_dim, - algo_name=playback_cfg.algo_log_name, + algo_name=guard_algo_name or playback_cfg.algo_log_name, ): - runner.load( - checkpoint_path, - load_cfg={ - "actor": True, - "critic": False, - "optimizer": False, - "iteration": False, - "rnd": False, - }, - ) + 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}") @@ -461,17 +576,22 @@ def create_rsl_rl_playback_session( action_mode=playback_cfg.action_mode, policy=policy, num_envs=playback_cfg.num_envs, + runner=runner, ) return session, policy_obs_mode, checkpoint_path -def _normalize_checkpoint_value(value: object) -> str | None: +def normalize_checkpoint_value(value: object) -> str | None: + """Normalize a raw checkpoint selector value; sentinel values map to ``None``.""" if value is None: return None text = str(value) return None if text in {"", "-1", "None", "null"} else text +_normalize_checkpoint_value = normalize_checkpoint_value + + def _cfg_checkpoint_value(cfg: Any) -> str | None: from omegaconf import OmegaConf @@ -648,6 +768,7 @@ def create_appo_playback_session( wrapped_env = selected_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) @@ -690,6 +811,7 @@ def create_appo_playback_session( action_mode=playback_cfg.action_mode, policy=policy, num_envs=playback_cfg.num_envs, + actor=actor, ), policy_obs_mode, checkpoint_path, @@ -715,6 +837,7 @@ def create_sac_playback_session( from unilab.training.offpolicy import ( default_device, extract_play_obs, + load_play_actor, resolve_play_actor_spec, resolve_play_obs_dims, ) @@ -788,10 +911,7 @@ def create_sac_playback_session( env_action_dim=action_dim, algo_name=algo_name, ): - actor.load_state_dict(checkpoint["actor"]) - if normalizer is not None and checkpoint.get("obs_normalizer"): - normalizer.load_state_dict(checkpoint["obs_normalizer"]) - normalizer.eval() + load_play_actor(algo_name, actor, normalizer, checkpoint) log(f"Loading {algo_name} checkpoint: {checkpoint_path}") log(f"Action mode: {playback_cfg.action_mode}") @@ -821,7 +941,6 @@ def _default_hora_distill_playback_deps(root_dir: str | Path) -> dict[str, Any]: ) from unilab.algos.hora.distill_config import apply_teacher_defaults from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - from unilab.base.backend import materialize_scene_visual_override from unilab.training import ( BackendAdapter, create_env, @@ -836,7 +955,7 @@ def _default_hora_distill_playback_deps(root_dir: str | Path) -> dict[str, Any]: cfg, root_dir=root_dir, algo_name="hora_distill", - scene_materializer=materialize_scene_visual_override, + scene_materializer=_scene_visual_materializer(), ).build_play_env_cfg_override(), "build_student_actor_and_normalizer": build_student_actor_and_normalizer, "cfg_with_checkpoint_runtime": cfg_with_checkpoint_runtime, @@ -1022,15 +1141,21 @@ def prepare_motion_overlay_selection( "KeyboardCommander", "MotionOverlaySelection", "OffPolicyPlaybackSession", + "PlayInteractiveArgs", "PlaybackControls", "PlaybackSession", "RslRlPlaybackConfig", "RslRlPlaybackSession", + "available_backends_for_task", + "build_play_backend_adapter", + "build_playback_config", "create_appo_playback_session", "create_hora_distill_playback_session", "create_rsl_rl_playback_session", "create_sac_playback_session", + "infer_checkpoint_actor_input_dim", "make_sim2sim_preflight", + "normalize_checkpoint_value", "prepare_motion_overlay_selection", "select_torch_device", ] diff --git a/tests/base/test_backend_imports.py b/tests/base/test_backend_imports.py index f7de2e03b..8e11229c8 100644 --- a/tests/base/test_backend_imports.py +++ b/tests/base/test_backend_imports.py @@ -9,10 +9,10 @@ _REPO_ROOT = Path(__file__).resolve().parents[2] _MATERIALIZER_CONSUMERS = ( "src/unilab/training/backend_adapter.py", + "src/unilab/visualization/interactive_playback.py", "scripts/train_rsl_rl.py", "scripts/train_him_ppo.py", "scripts/train_hora_distill.py", - "scripts/play_interactive.py", "scripts/manip_loco/benchmark_site_jacobian.py", ) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 83ba1398d..ba8419c0a 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -1876,22 +1876,28 @@ def test_resolve_checkpoint_empty_run_dir(tmp_path): def test_offpolicy_extract_reset_obs_handles_two_tuple(): + from unilab.training.offpolicy import extract_reset_obs + obs = {"obs": "value"} - result = _offpolicy().extract_reset_obs((obs, {"info": 1})) + result = extract_reset_obs((obs, {"info": 1})) assert result is obs def test_offpolicy_extract_reset_obs_rejects_three_tuple(): + from unilab.training.offpolicy import extract_reset_obs + obs = {"obs": "value"} with pytest.raises(ValueError, match="Unexpected env.reset return format"): - _offpolicy().extract_reset_obs(("ignored", obs, {"info": 1})) + extract_reset_obs(("ignored", obs, {"info": 1})) def test_offpolicy_resolve_play_obs_dim_ignores_critic(): - obs_dim = _offpolicy().resolve_play_obs_dim({"obs": 98, "critic": 101}) + from unilab.training.offpolicy import resolve_play_obs_dim + + obs_dim = resolve_play_obs_dim({"obs": 98, "critic": 101}) assert obs_dim == 98 @@ -1899,12 +1905,14 @@ def test_offpolicy_resolve_play_obs_dim_ignores_critic(): def test_offpolicy_extract_play_obs_uses_obs_group_only(): import numpy as np + from unilab.training.offpolicy import extract_play_obs + obs = { "obs": np.ones((2, 98), dtype=np.float32), "critic": np.full((2, 101), 2.0, dtype=np.float32), } - play_obs = _offpolicy().extract_play_obs(obs) + play_obs = extract_play_obs(obs) assert play_obs.shape == (2, 98) assert np.allclose(play_obs, 1.0) @@ -2141,7 +2149,10 @@ def step(self, actions): self.state = type( "State", (), - {"obs": {"obs": np.ones((batch, 4), dtype=np.float32)}}, + { + "obs": {"obs": np.ones((batch, 4), dtype=np.float32)}, + "info": {}, + }, )() captured["actions_shape"] = actions.shape return self.state @@ -2158,13 +2169,19 @@ def run_playback_mode(self, **kwargs): monkeypatch.setattr(mod, "build_offpolicy_env_cfg_override", lambda algo_name, cfg: {}) monkeypatch.setattr(mod, "default_device", lambda torch_module, preferred=None: "cpu") monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: FakeEnv()) - monkeypatch.setattr(mod, "resolve_play_obs_dim", lambda obs_groups_spec: 4) - monkeypatch.setattr(mod, "extract_play_obs", lambda obs_dict: obs_dict["obs"]) monkeypatch.setattr( mod, "resolve_checkpoint_path", lambda *args, **kwargs: (str(checkpoint), str(run_dir)), ) + + import unilab.training.run as training_run + + monkeypatch.setattr( + training_run, + "resolve_offpolicy_checkpoint_path", + lambda *args, **kwargs: (str(checkpoint), str(run_dir)), + ) monkeypatch.setattr( torch.onnx, "export", @@ -2295,6 +2312,14 @@ def run_playback_mode(self, **kwargs): lambda *args, **kwargs: (str(checkpoint), str(run_dir)), ) + import unilab.training.run as training_run + + monkeypatch.setattr( + training_run, + "resolve_offpolicy_checkpoint_path", + lambda *args, **kwargs: (str(checkpoint), str(run_dir)), + ) + import unilab.algos.common.actor_factory as actor_factory def fake_build_actor(algo_type, obs_dim, action_dim, hidden_dim, use_layer_norm, device, **kw): @@ -2896,6 +2921,8 @@ def test_train_rsl_rl_motrix_auto_play_is_interactive( class FakeEnv: def __init__(self): self.cfg = type("Cfg", (), {"render_spacing": 2.5, "render_offset_mode": "zero"})() + self.obs_groups_spec = {"obs": 1} + self.action_space = type("Space", (), {"shape": (1,)})() def run_playback_mode(self, **kwargs): assert kwargs["play_render_mode"] == "auto" @@ -2921,9 +2948,10 @@ def run_playback_mode(self, **kwargs): return None class FakeWrapper: - def __init__(self, env, device): + def __init__(self, env, device, policy_obs_mode="flat"): self.env = env self.device = device + self.policy_obs_mode = policy_obs_mode def reset(self): return 0, {} @@ -2986,6 +3014,8 @@ def test_train_rsl_rl_record_play_uses_backend_plan( class FakeEnv: def __init__(self): self.cfg = type("Cfg", (), {"render_spacing": 1.0, "render_offset_mode": "grid"})() + self.obs_groups_spec = {"obs": 1} + self.action_space = type("Space", (), {"shape": (1,)})() def run_playback_mode(self, **kwargs): assert kwargs["play_render_mode"] == "record" @@ -3011,9 +3041,10 @@ def run_playback_mode(self, **kwargs): return str(plan.output_video) class FakeWrapper: - def __init__(self, env, device): + def __init__(self, env, device, policy_obs_mode="flat"): self.env = env self.device = device + self.policy_obs_mode = policy_obs_mode def reset(self): return 0, {} @@ -3028,9 +3059,9 @@ def __init__(self, wrapped_env, train_cfg, log_dir, device): self.log_dir = log_dir self.device = device - def load(self, path, map_location=None): + def load(self, path, **kwargs): self.loaded_path = path - self.map_location = map_location + self.load_kwargs = kwargs def get_inference_policy(self, device): return lambda obs: obs @@ -3325,3 +3356,387 @@ def _fail_bootstrap() -> None: with pytest.raises(RuntimeError, match="bootstrap failed"): spec.loader.exec_module(mod) # type: ignore[union-attr] + + +# --------------------------------------------------------------------------- +# Unified play entrypoints — shared playback session factories (issue #1242) +# --------------------------------------------------------------------------- + + +def _him_ppo_cfg(overrides=None): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo_him"), version_base="1.3"): + return compose( + "config", + overrides=["task=go2_arm_manip_loco/mujoco", *(overrides or [])], + ) + + +def _train_him_ppo(): + return _load_script("train_him_ppo") + + +class _FakePlaybackEnv: + """Env stand-in driving one initialize/step cycle through run_playback_mode.""" + + def __init__(self, video_path: str): + self.cfg = types.SimpleNamespace(render_spacing=1.0, render_offset_mode="grid") + self.video_path = video_path + self.captured: dict[str, Any] = {} + + def run_playback_mode(self, **kwargs: Any) -> str: + self.captured.update(kwargs) + self.captured["init_obs"] = kwargs["initialize"]() + self.captured["next_obs"] = kwargs["step"](self.captured["init_obs"]) + return self.video_path + + +def test_train_rsl_rl_play_uses_shared_playback_session_factory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + mod = _train_rsl_rl(monkeypatch) + cfg = _ppo_cfg( + [ + "task=go1_joystick_flat/mujoco", + "training.play_only=true", + "training.play_render_mode=record", + "training.play_steps=5", + ] + ) + run_dir = tmp_path / "run" + run_dir.mkdir() + checkpoint = run_dir / "model_37.pt" + mod.torch.save({"actor_state_dict": {}}, checkpoint) + captured: dict[str, Any] = {} + + class FakeSession: + def __init__(self): + self.env = _FakePlaybackEnv(str(run_dir / "play_video.mp4")) + self.runner = object() + self.reset_calls = 0 + self.step_calls = 0 + + def reset(self): + self.reset_calls += 1 + return "obs_0" + + def step_once(self): + self.step_calls += 1 + return "obs_1" + + fake_session = FakeSession() + sentinel_wrapper_cls = type("SentinelWrapper", (), {}) + + def fake_create_session(**kwargs: Any): + captured["factory_kwargs"] = kwargs + return fake_session, "flat", str(checkpoint) + + monkeypatch.setattr(mod, "EXPORT_POLICY", False, raising=False) + monkeypatch.setattr(mod, "parse_checkpoint_path", lambda *args, **kwargs: (checkpoint, run_dir)) + monkeypatch.setattr(mod, "_resolve_ppo_wrapper_cls", lambda rl_cfg: sentinel_wrapper_cls) + monkeypatch.setattr(mod, "create_rsl_rl_playback_session", fake_create_session) + + result = mod.play_rsl_rl(cfg, device="cpu") + + assert result == str(run_dir / "play_video.mp4") + factory_kwargs = captured["factory_kwargs"] + playback_cfg = factory_kwargs["playback_cfg"] + assert playback_cfg.task == cfg.training.task_name + assert playback_cfg.action_mode == "policy" + assert playback_cfg.policy_obs_mode == "flat" + assert playback_cfg.algo_log_name == cfg.algo.algo_log_name + assert playback_cfg.num_envs == cfg.training.play_env_num + assert factory_kwargs["device"] == "cpu" + assert factory_kwargs["root_dir"] == mod.ROOT_DIR + assert factory_kwargs["wrapper_cls"] is sentinel_wrapper_cls + assert factory_kwargs["runner_cls"] is mod.OnPolicyRunner + assert factory_kwargs["guard_algo_name"] == "ppo" + assert factory_kwargs.get("runner_loader") is None + assert factory_kwargs["checkpoint_resolver"]() == str(checkpoint) + assert callable(factory_kwargs["sim2sim_preflight"]) + assert fake_session.reset_calls == 1 + assert fake_session.step_calls == 1 + env_captured = fake_session.env.captured + assert env_captured["play_render_mode"] == "record" + assert env_captured["play_steps"] == 5 + assert env_captured["init_obs"] == "obs_0" + assert env_captured["next_obs"] == "obs_1" + + +def test_train_him_ppo_play_missing_checkpoint_returns_none_without_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str] +): + mod = _train_him_ppo() + cfg = _him_ppo_cfg(["training.play_only=true"]) + + monkeypatch.setattr(mod, "parse_checkpoint_path", lambda *args, **kwargs: (None, None)) + monkeypatch.setattr( + mod, + "create_env", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("play_him_ppo should not create an env before checkpoint resolution") + ), + ) + + result = mod.play_him_ppo(cfg, device="cpu") + + assert result is None + assert "Could not resolve a checkpoint for play mode." in capsys.readouterr().out + + +def test_train_him_ppo_play_uses_shared_playback_session_factory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + mod = _train_him_ppo() + cfg = _him_ppo_cfg(["training.play_only=true"]) + run_dir = tmp_path / "run" + run_dir.mkdir() + checkpoint = run_dir / "model_37.pt" + mod.torch.save({"actor_state_dict": {}}, checkpoint) + captured: dict[str, Any] = {} + + class FakeSession: + def __init__(self): + self.env = types.SimpleNamespace( + cfg=types.SimpleNamespace(render_spacing=1.0), + ) + self.runner = object() + self.policy = lambda obs: obs + self.reset_calls = 0 + self.step_calls = 0 + + def reset(self): + self.reset_calls += 1 + return {"actor": "obs_0"} + + def step_once(self): + self.step_calls += 1 + return {"actor": "obs_1"} + + fake_session = FakeSession() + + def fake_create_session(**kwargs: Any): + captured["factory_kwargs"] = kwargs + return fake_session, "actor", str(checkpoint) + + def fake_render_play_mode(env, **kwargs: Any): + captured["render_kwargs"] = kwargs + captured["init_obs"] = kwargs["initialize"]() + captured["next_obs"] = kwargs["step"](captured["init_obs"]) + + monkeypatch.setattr(mod, "EXPORT_POLICY", False, raising=False) + monkeypatch.setattr(mod, "parse_checkpoint_path", lambda *args, **kwargs: (checkpoint, run_dir)) + monkeypatch.setattr(mod, "create_rsl_rl_playback_session", fake_create_session) + monkeypatch.setattr(mod, "render_play_mode", fake_render_play_mode) + + result = mod.play_him_ppo(cfg, device="cpu") + + assert result == str(run_dir / "play_video.mp4") + factory_kwargs = captured["factory_kwargs"] + playback_cfg = factory_kwargs["playback_cfg"] + assert playback_cfg.task == cfg.training.task_name + assert playback_cfg.action_mode == "policy" + assert playback_cfg.num_envs == cfg.training.play_env_num + assert factory_kwargs["device"] == "cpu" + assert factory_kwargs["wrapper_cls"] is mod.RslRlVecEnvWrapper + assert factory_kwargs["runner_cls"] is mod.HIMOnPolicyRunner + assert factory_kwargs["guard_algo_name"] == "him_ppo" + assert callable(factory_kwargs["runner_loader"]) + assert factory_kwargs["checkpoint_resolver"]() == str(checkpoint) + assert callable(factory_kwargs["sim2sim_preflight"]) + assert fake_session.reset_calls == 1 + assert fake_session.step_calls == 1 + assert captured["init_obs"] == "obs_0" + assert captured["next_obs"] == "obs_1" + assert captured["render_kwargs"]["output_video"] == run_dir / "play_video.mp4" + + +def test_play_appo_missing_checkpoint_returns_none_without_env( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +): + mod = _train_appo() + cfg = _appo_cfg(["task=g1_walk_flat/mujoco", "training.play_only=true"]) + + monkeypatch.setattr( + mod, + "create_env", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("play_appo should not create an env before checkpoint resolution") + ), + ) + + result = mod.play_appo(cfg, {}, resolve_checkpoint_path=lambda _cfg: (None, None)) + + assert result is None + assert "Could not find run to load." in capsys.readouterr().out + + +def test_play_appo_uses_shared_playback_session_factory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + import torch + + mod = _train_appo() + cfg = _appo_cfg( + [ + "task=g1_walk_flat/mujoco", + "training.play_only=true", + "training.play_render_mode=record", + ] + ) + run_dir = tmp_path / "run" + run_dir.mkdir() + checkpoint = run_dir / "model_37.pt" + checkpoint.write_bytes(b"checkpoint") + captured: dict[str, Any] = {} + + class FakeActor: + def __init__(self): + self.mlp = torch.nn.Linear(4, 2) + + class FakeSession: + def __init__(self): + self.env = _FakePlaybackEnv(str(run_dir / "play_video.mp4")) + self.actor = FakeActor() + self.wrapped_env = types.SimpleNamespace(num_obs=4) + self.reset_calls = 0 + self.step_calls = 0 + + def reset(self): + self.reset_calls += 1 + return "obs_0" + + def step_once(self): + self.step_calls += 1 + return "obs_1" + + fake_session = FakeSession() + rl_cfg: dict[str, Any] = {"seed": 1} + + def fake_create_session(**kwargs: Any): + captured["factory_kwargs"] = kwargs + return fake_session, "flat", str(checkpoint) + + monkeypatch.setattr(mod, "create_appo_playback_session", fake_create_session) + monkeypatch.setattr( + mod, "export_policy_onnx", lambda *args, **kwargs: captured.setdefault("onnx_export", args) + ) + monkeypatch.setattr(mod, "verify_policy_onnx", lambda *args, **kwargs: None) + + result = mod.play_appo( + cfg, + rl_cfg, + resolve_checkpoint_path=lambda _cfg: (str(checkpoint), str(run_dir)), + ) + + assert result == str(run_dir / "play_video.mp4") + factory_kwargs = captured["factory_kwargs"] + playback_cfg = factory_kwargs["playback_cfg"] + assert playback_cfg.task == cfg.training.task_name + assert playback_cfg.action_mode == "policy" + assert playback_cfg.algo_log_name == cfg.algo.algo_log_name + assert playback_cfg.num_envs == cfg.training.play_env_num + assert factory_kwargs["cfg"] is cfg + assert factory_kwargs["rl_cfg"] is rl_cfg + assert factory_kwargs["root_dir"] == mod.ROOT_DIR + assert factory_kwargs["wrapper_cls"] is mod.RslRlVecEnvWrapper + assert "onnx_export" in captured + assert fake_session.reset_calls == 1 + assert fake_session.step_calls == 1 + env_captured = fake_session.env.captured + assert env_captured["play_render_mode"] == "record" + assert env_captured["init_obs"] == "obs_0" + assert env_captured["next_obs"] == "obs_1" + + +def test_play_offpolicy_missing_checkpoint_returns_none_without_env( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +): + mod = _offpolicy() + cfg = _offpolicy_cfg(["task=g1_walk_flat/mujoco", "training.play_only=true"]) + + monkeypatch.setattr(mod, "resolve_checkpoint_path", lambda *args, **kwargs: (None, None)) + monkeypatch.setattr( + mod, + "create_env", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("play_offpolicy should not create an env before checkpoint resolution") + ), + ) + + result = mod.play_offpolicy("sac", cfg) + + assert result is None + assert "Could not find checkpoint." in capsys.readouterr().out + + +def test_play_offpolicy_uses_shared_playback_session_factory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "task=g1_walk_flat/mujoco", + "training.play_only=true", + "training.play_render_mode=record", + "training.export_onnx=false", + ] + ) + run_dir = tmp_path / "run" + run_dir.mkdir() + checkpoint = run_dir / "model_5000.pt" + checkpoint.write_bytes(b"checkpoint") + captured: dict[str, Any] = {} + + class FakeSession: + def __init__(self): + self.env = _FakePlaybackEnv(str(run_dir / "play_video.mp4")) + self.actor = object() + self.normalizer = None + self.actor_algo_type = "sac" + self.reset_calls = 0 + self.step_calls = 0 + + def reset(self): + self.reset_calls += 1 + return "obs_0" + + def step_once(self): + self.step_calls += 1 + return "obs_1" + + fake_session = FakeSession() + + def fake_create_session(**kwargs: Any): + captured["factory_kwargs"] = kwargs + return fake_session, "actor", str(checkpoint) + + monkeypatch.setattr(mod, "default_device", lambda torch_module, preferred=None: "cpu") + monkeypatch.setattr( + mod, + "resolve_checkpoint_path", + lambda *args, **kwargs: (str(checkpoint), str(run_dir)), + ) + monkeypatch.setattr(mod, "create_sac_playback_session", fake_create_session) + + result = mod.play_offpolicy("sac", cfg) + + assert result == str(run_dir / "play_video.mp4") + factory_kwargs = captured["factory_kwargs"] + playback_cfg = factory_kwargs["playback_cfg"] + assert playback_cfg.task == cfg.training.task_name + assert playback_cfg.action_mode == "policy" + assert playback_cfg.policy_obs_mode == "actor" + assert playback_cfg.algo_log_name == cfg.algo.algo_log_name + assert playback_cfg.num_envs == cfg.training.play_env_num + assert factory_kwargs["algo_name"] == "sac" + assert factory_kwargs["device"] == "cpu" + assert factory_kwargs["cfg"] is cfg + assert factory_kwargs["root_dir"] == mod.ROOT_DIR + assert callable(factory_kwargs["env_factory"]) + assert fake_session.reset_calls == 1 + assert fake_session.step_calls == 1 + env_captured = fake_session.env.captured + assert env_captured["play_render_mode"] == "record" + assert env_captured["init_obs"] == "obs_0" + assert env_captured["next_obs"] == "obs_1" diff --git a/tests/visualization/test_interactive_playback.py b/tests/visualization/test_interactive_playback.py index da62e95d7..576069ff1 100644 --- a/tests/visualization/test_interactive_playback.py +++ b/tests/visualization/test_interactive_playback.py @@ -767,6 +767,7 @@ def forward(self, obs): session.reset() assert session.advance(PlaybackControls()) is True + assert isinstance(session.actor, FakeActor) assert policy_obs_mode == "actor" assert resolved_checkpoint == str(checkpoint) assert captured["wrapper_cls"] == "hora" @@ -1038,3 +1039,362 @@ def fake_load_distilled_checkpoint(actor, hist_normalizer, checkpoint_path, *, d assert captured["policy_obs_mode"] == "actor" assert captured["loaded_checkpoint"] == checkpoint torch.testing.assert_close(captured["actions"], torch.ones((1, 2))) + + +def test_infer_checkpoint_actor_input_dim_mlp_key(tmp_path: Path) -> None: + from unilab.visualization.interactive_playback import infer_checkpoint_actor_input_dim + + checkpoint = tmp_path / "model_10.pt" + torch.save({"actor_state_dict": {"mlp.0.weight": torch.zeros((8, 42))}}, checkpoint) + + assert infer_checkpoint_actor_input_dim(str(checkpoint)) == 42 + + +def test_infer_checkpoint_actor_input_dim_actor_prefixed_key(tmp_path: Path) -> None: + from unilab.visualization.interactive_playback import infer_checkpoint_actor_input_dim + + checkpoint = tmp_path / "model_10.pt" + torch.save( + {"actor_state_dict": {"actor.mlp.0.weight": torch.zeros((8, 17))}}, + checkpoint, + ) + + assert infer_checkpoint_actor_input_dim(str(checkpoint)) == 17 + + +def test_infer_checkpoint_actor_input_dim_generic_first_layer_key(tmp_path: Path) -> None: + from unilab.visualization.interactive_playback import infer_checkpoint_actor_input_dim + + checkpoint = tmp_path / "model_10.pt" + torch.save( + {"actor_state_dict": {"encoder.0.weight": torch.zeros((4, 23))}}, + checkpoint, + ) + + assert infer_checkpoint_actor_input_dim(str(checkpoint)) == 23 + + +def test_infer_checkpoint_actor_input_dim_returns_none_when_undetectable( + tmp_path: Path, +) -> None: + from unilab.visualization.interactive_playback import infer_checkpoint_actor_input_dim + + not_a_dict = tmp_path / "model_list.pt" + torch.save({"actor_state_dict": ["not", "a", "dict"]}, not_a_dict) + missing = tmp_path / "model_missing.pt" + torch.save({"model_state_dict": {}}, missing) + no_matching_key = tmp_path / "model_other.pt" + torch.save( + {"actor_state_dict": {"mlp.2.weight": torch.zeros((8, 8))}}, + no_matching_key, + ) + + assert infer_checkpoint_actor_input_dim(str(not_a_dict)) is None + assert infer_checkpoint_actor_input_dim(str(missing)) is None + assert infer_checkpoint_actor_input_dim(str(no_matching_key)) is None + + +@pytest.mark.parametrize("value", [None, "", "-1", "None", "null"]) +def test_normalize_checkpoint_value_maps_sentinels_to_none(value: object) -> None: + from unilab.visualization.interactive_playback import normalize_checkpoint_value + + assert normalize_checkpoint_value(value) is None + + +@pytest.mark.parametrize( + ("value", "expected"), + [("12", "12"), (12, "12"), ("run_2024", "run_2024"), ("/abs/model_5.pt", "/abs/model_5.pt")], +) +def test_normalize_checkpoint_value_keeps_real_values(value: object, expected: str) -> None: + from unilab.visualization.interactive_playback import normalize_checkpoint_value + + assert normalize_checkpoint_value(value) == expected + + +def _play_interactive_args(**overrides: Any) -> Any: + from unilab.visualization.interactive_playback import PlayInteractiveArgs + + defaults: dict[str, Any] = { + "task": "MyTask", + "load_run": "-1", + "checkpoint": None, + "action_mode": "policy", + "policy_obs_mode": "auto", + "algo_log_name": "rsl_rl_ppo", + "log_root": None, + "show_target_bodies": False, + "show_reward_debug": False, + "target_show_axes": False, + "target_body_names": "", + "target_max_bodies": 32, + "target_marker_radius": 0.05, + "target_axis_length": 0.2, + "target_marker_alpha": 0.7, + "reward_debug_show_velocity": False, + "reward_debug_lin_vel_scale": 1.0, + "reward_debug_ang_vel_scale": 1.0, + "reward_debug_show_connectors": False, + "reward_debug_show_global_anchor": False, + "camera_follow_body": False, + "camera_focus_body_name": "", + "camera_height_offset": 0.0, + "camera_distance": None, + "camera_elevation": None, + "camera_azimuth": None, + "use_env_visual_model": False, + "speed": 1.0, + "start_paused": False, + } + defaults.update(overrides) + return PlayInteractiveArgs(**defaults) + + +def test_build_playback_config_maps_play_interactive_args() -> None: + from unilab.visualization.interactive_playback import build_playback_config + + args = _play_interactive_args( + task="OtherTask", + load_run="run_1", + checkpoint="12", + action_mode="random", + policy_obs_mode="flat", + algo_log_name="custom_ppo", + log_root="/tmp/logs", + speed=2.5, + start_paused=True, + ) + + playback_cfg = build_playback_config(args, num_envs=3) + + assert playback_cfg == RslRlPlaybackConfig( + task="OtherTask", + load_run="run_1", + checkpoint="12", + action_mode="random", + policy_obs_mode="flat", + algo_log_name="custom_ppo", + log_root="/tmp/logs", + num_envs=3, + speed=2.5, + start_paused=True, + ) + + +def test_available_backends_for_task_reads_registry(monkeypatch: pytest.MonkeyPatch) -> None: + from unilab.base import registry + from unilab.visualization.interactive_playback import available_backends_for_task + + monkeypatch.setattr( + registry, + "list_registered_envs", + lambda: { + "KnownTask": {"available_backends": ["mujoco", "motrix"]}, + "BadTask": {"available_backends": "mujoco"}, + }, + ) + + assert available_backends_for_task("KnownTask") == ("mujoco", "motrix") + assert available_backends_for_task("UnknownTask") == () + assert available_backends_for_task("BadTask") == () + + +def test_build_play_backend_adapter_injects_root_dir_and_materializer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import unilab.base.backend as backend_pkg + import unilab.training as training + from unilab.visualization.interactive_playback import build_play_backend_adapter + + captured: dict[str, Any] = {} + sentinel_materializer = object() + + class FakeBackendAdapter: + def __init__(self, cfg: Any, **kwargs: Any) -> None: + captured["cfg"] = cfg + captured.update(kwargs) + + monkeypatch.setattr(training, "BackendAdapter", FakeBackendAdapter) + monkeypatch.setattr( + backend_pkg, + "materialize_scene_visual_override", + sentinel_materializer, + ) + + cfg = SimpleNamespace(training=SimpleNamespace(task_name="Task")) + adapter = build_play_backend_adapter(cfg, root_dir="/repo", algo_name="appo") + + assert isinstance(adapter, FakeBackendAdapter) + assert captured == { + "cfg": cfg, + "root_dir": "/repo", + "algo_name": "appo", + "scene_materializer": sentinel_materializer, + } + + +def test_create_rsl_rl_playback_session_uses_runner_loader_and_exposes_runner( + tmp_path: Path, +) -> None: + run_dir = tmp_path / "run_1" + run_dir.mkdir() + checkpoint = run_dir / "model_10.pt" + torch.save({"actor_state_dict": {}}, checkpoint) + captured: dict[str, Any] = {} + + class Runner: + def __init__(self, wrapped_env, train_cfg, log_dir, device): + pass + + def load(self, checkpoint, load_cfg): + raise AssertionError("runner.load must be bypassed when runner_loader is injected") + + def get_inference_policy(self, *, device): + return lambda obs: torch.ones((1, 2)) + + def runner_loader(runner, path): + captured["loader_runner"] = runner + captured["loader_path"] = path + + kwargs = _rsl_rl_session_kwargs(tmp_path) + kwargs["checkpoint_resolver"] = lambda *args: str(checkpoint) + kwargs["runner_cls"] = Runner + kwargs["runner_loader"] = runner_loader + + session, _mode, resolved = create_rsl_rl_playback_session(**kwargs) + + assert resolved == str(checkpoint) + assert isinstance(captured["loader_runner"], Runner) + assert captured["loader_path"] == str(checkpoint) + assert session.runner is captured["loader_runner"] + + +def test_create_rsl_rl_playback_session_forwards_guard_algo_name( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import unilab.visualization.interactive_playback as interactive_playback + + run_dir = tmp_path / "run_1" + run_dir.mkdir() + checkpoint = run_dir / "model_10.pt" + torch.save({"actor_state_dict": {}}, checkpoint) + captured: dict[str, Any] = {} + + class Runner: + def __init__(self, wrapped_env, train_cfg, log_dir, device): + pass + + def load(self, checkpoint, load_cfg): + pass + + def get_inference_policy(self, *, device): + return lambda obs: torch.ones((1, 2)) + + import contextlib + + @contextlib.contextmanager + def fake_dim_guard(**kwargs): + captured["dim_guard"] = kwargs + yield + + monkeypatch.setattr(interactive_playback, "policy_load_dim_guard", fake_dim_guard) + + kwargs = _rsl_rl_session_kwargs(tmp_path) + kwargs["checkpoint_resolver"] = lambda *args: str(checkpoint) + kwargs["runner_cls"] = Runner + kwargs["guard_algo_name"] = "him_ppo" + + create_rsl_rl_playback_session(**kwargs) + + assert captured["dim_guard"] == { + "env_obs_dim": 5, + "env_action_dim": 2, + "algo_name": "him_ppo", + } + + +def test_create_sac_playback_session_td3_load_filters_noise_scales( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from omegaconf import OmegaConf + + import unilab.algos.common.actor_factory as actor_factory + import unilab.training.run as training_run + + checkpoint = tmp_path / "model_10.pt" + torch.save( + {"actor": {"weight": torch.ones(1), "noise_scales": torch.zeros(1)}}, + checkpoint, + ) + captured: dict[str, Any] = {} + + class FakeActor: + def eval(self): + return self + + def load_state_dict(self, state_dict, strict=True): + captured["actor_load"] = (state_dict, strict) + + class FakeEnv: + num_envs = 1 + obs_groups_spec = {"obs": 3, "critic": 5} + action_space = SimpleNamespace( + shape=(2,), + low=np.full((2,), -1.0), + high=np.full((2,), 1.0), + ) + state = SimpleNamespace(info={}) + + def get_physics_state_snapshot(self): + return np.zeros((1, 4), dtype=np.float32) + + cfg = OmegaConf.create( + { + "training": {"task_name": "Task", "device": None}, + "algo": { + "algo_log_name": "td3", + "load_run": "run", + "actor_hidden_dim": 16, + "use_layer_norm": False, + }, + } + ) + + def fake_build_actor(algo_type, *args, **kwargs): + captured["build_actor_algo_type"] = algo_type + return FakeActor() + + monkeypatch.setattr(actor_factory, "build_actor", fake_build_actor) + monkeypatch.setattr( + training_run, + "resolve_offpolicy_checkpoint_path", + lambda *args, **kwargs: (str(checkpoint), str(tmp_path)), + ) + + session, policy_obs_mode, resolved = create_sac_playback_session( + playback_cfg=RslRlPlaybackConfig( + task="Task", + load_run="run", + checkpoint=None, + action_mode="policy", + policy_obs_mode="actor", + algo_log_name="td3", + log_root=None, + ), + cfg=cfg, + env_factory=lambda num_envs: FakeEnv(), + root_dir=tmp_path, + device="cpu", + algo_name="td3", + log=lambda message: None, + ) + + assert resolved == str(checkpoint) + assert policy_obs_mode == "actor" + assert captured["build_actor_algo_type"] == "td3" + assert isinstance(session.actor, FakeActor) + assert session.actor_algo_type == "td3" + loaded_state, strict = captured["actor_load"] + assert set(loaded_state.keys()) == {"weight"} + assert strict is False