From 8d019e15653cd0cfd3041fc70307c708b82a61cb Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Fri, 10 Apr 2026 01:20:03 +0800 Subject: [PATCH] refactor: extract shared training entrypoint helpers --- scripts/train_appo.py | 173 +++++-------- scripts/train_mlx_ppo.py | 200 +++++---------- scripts/train_offpolicy.py | 267 ++++++--------------- scripts/train_rsl_rl.py | 307 ++++-------------------- src/unilab/training/__init__.py | 27 +++ src/unilab/training/backend_adapter.py | 161 +++++++++++++ src/unilab/training/common.py | 225 +++++++++++++++++ tests/training/test_training_helpers.py | 102 ++++++++ 8 files changed, 755 insertions(+), 707 deletions(-) create mode 100644 src/unilab/training/__init__.py create mode 100644 src/unilab/training/backend_adapter.py create mode 100644 src/unilab/training/common.py create mode 100644 tests/training/test_training_helpers.py diff --git a/scripts/train_appo.py b/scripts/train_appo.py index c3e6fb12c..a5a1eb289 100644 --- a/scripts/train_appo.py +++ b/scripts/train_appo.py @@ -3,9 +3,7 @@ from __future__ import annotations import datetime -import importlib import os -import pkgutil import sys from pathlib import Path @@ -16,27 +14,10 @@ ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) +from unilab.training import create_env, ensure_registries, get_log_root, render_play_mode from unilab.utils.experiment_tracking import ExperimentTracker -def ensure_registries(): - for pkg_name in ( - "unilab.envs.locomotion", - "unilab.envs.manipulation", - "unilab.envs.motion_tracking", - ): - try: - package = importlib.import_module(pkg_name) - if hasattr(package, "__path__"): - for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): - try: - importlib.import_module(name) - except Exception: - pass - except ImportError: - pass - - def build_appo_runner_kwargs( cfg: DictConfig, env_cfg_override: dict | None, @@ -68,79 +49,52 @@ def run_motrix_play_loop( play_env_num: int, num_steps: int | None = None, ) -> None: - import time - import numpy as np from tensordict import TensorDict if env.state is None: env.init_state() - env._backend.init_renderer() - - env_indices = np.arange(play_env_num, dtype=np.int32) - obs_out, _ = env.reset(env_indices) - obs_np = np.asarray(obs_out["obs"], dtype=np.float32) - - last_render_time = time.perf_counter() - render_dt = 1.0 / 60.0 - steps_run = 0 with torch.inference_mode(): - while num_steps is None or steps_run < num_steps: - obs_torch = torch.from_numpy(obs_np).to(device) - td = TensorDict({"policy": obs_torch}, batch_size=play_env_num) - actions_torch = actor(td) - actions_np = actions_torch.cpu().numpy().astype(np.float32) - state = env.step(actions_np) - obs_np = np.asarray(state.obs["obs"], dtype=np.float32) - - current_time = time.perf_counter() - elapsed = current_time - last_render_time - if elapsed < render_dt: - time.sleep(render_dt - elapsed) - last_render_time = time.perf_counter() - - env._backend.render() - steps_run += 1 + render_play_mode( + env, + sim_backend="motrix", + num_steps=num_steps, + initialize=lambda: np.asarray( + env.reset(np.arange(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=play_env_num + ) + ) + .cpu() + .numpy() + .astype(np.float32) + ).obs["obs"], + dtype=np.float32, + ), + ) def resolve_appo_checkpoint_path( base_log_dir: str | Path, load_run: str, ) -> tuple[str | None, str | None]: - base_log_dir = str(base_log_dir) + from unilab.training import resolve_checkpoint_path - if load_run == "-1": - if not os.path.exists(base_log_dir): - return None, None - all_runs = sorted( - [d for d in os.listdir(base_log_dir) if os.path.isdir(os.path.join(base_log_dir, d))] - ) - if not all_runs: - return None, None - candidate_dir = os.path.join(base_log_dir, all_runs[-1]) - elif os.path.isdir(load_run): - candidate_dir = load_run - elif os.path.isfile(load_run): - return load_run, os.path.dirname(load_run) - else: - candidate_dir = os.path.join(base_log_dir, load_run) - if not os.path.isdir(candidate_dir): - return None, None - - model_files = [ - f for f in os.listdir(candidate_dir) if f.startswith("model_") and f.endswith(".pt") - ] - if not model_files: - return None, None - - model_files.sort(key=lambda x: int(x.split("_")[1].split(".")[0])) - return os.path.join(candidate_dir, model_files[-1]), candidate_dir + checkpoint_path, checkpoint_dir = resolve_checkpoint_path(base_log_dir, load_run, suffix=".pt") + return ( + str(checkpoint_path) if checkpoint_path is not None else None, + str(checkpoint_dir) if checkpoint_dir is not None else None, + ) def _get_log_root(cfg: DictConfig) -> str: - """Get log root directory from algo_log_name config.""" - return str(Path(ROOT_DIR) / "logs" / cfg.algo.algo_log_name) + return str(get_log_root(ROOT_DIR, cfg)) def play_appo(cfg: DictConfig, rl_cfg: dict) -> str | None: @@ -149,7 +103,6 @@ def play_appo(cfg: DictConfig, rl_cfg: dict) -> str | None: from rsl_rl.utils import resolve_callable from tensordict import TensorDict - from unilab.base import registry from unilab.utils.reward_utils import extract_reward_config from unilab.utils.rsl_rl_compat import convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5 @@ -164,10 +117,9 @@ def play_appo(cfg: DictConfig, rl_cfg: dict) -> str | None: ) print(f"Using device for play: {device}") - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=cfg.training.play_env_num, - sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) from unilab.utils.obs_utils import get_obs_dims @@ -231,47 +183,46 @@ def play_appo(cfg: DictConfig, rl_cfg: dict) -> str | None: raise return None - import mediapy as media - - from unilab.utils import render_many - output_video = os.path.join(load_path_dir, "play_video.mp4") print(f"Rendering video to {output_video}...") if env.state is None: env.init_state() - env_indices = np.arange(cfg.training.play_env_num, dtype=np.int32) - obs_out, _ = env.reset(env_indices) - obs_np = np.asarray(obs_out["obs"], dtype=np.float32) - - state_list = [] num_steps = int(getattr(cfg.training, "play_steps", 1000)) print("Collecting physics states...") with torch.inference_mode(): - for _ in range(num_steps): - obs_torch = torch.from_numpy(obs_np).to(device) - td = TensorDict({"policy": obs_torch}, batch_size=cfg.training.play_env_num) - actions_torch = actor(td) - actions_np = actions_torch.cpu().numpy().astype(np.float32) - state = env.step(actions_np) - obs_np = np.asarray(state.obs["obs"], dtype=np.float32) - state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy()) - - print("Rendering frames...") - frames = render_many.render_states_get_frames( - state_list, - env.cfg.model_file, - width=1280, - height=720, - camera_id=-1, - cam_distance=cfg.training.cam_distance, - cam_elevation=cfg.training.cam_elevation, - cam_azimuth=cfg.training.cam_azimuth, - ) - + render_play_mode( + env, + sim_backend=cfg.training.sim_backend, + num_steps=num_steps, + output_video=output_video, + 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, + ), + frame_state_getter=lambda: env._backend.get_physics_state(), + camera_kwargs={ + "cam_distance": cfg.training.cam_distance, + "cam_elevation": cfg.training.cam_elevation, + "cam_azimuth": cfg.training.cam_azimuth, + }, + ) print(f"Saving video to {output_video} with mediapy...") - media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt)) print("Done.") return output_video diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py index 22c156378..7f85b033f 100644 --- a/scripts/train_mlx_ppo.py +++ b/scripts/train_mlx_ppo.py @@ -5,11 +5,9 @@ from __future__ import annotations import datetime -import importlib import math import os import pickle -import pkgutil import statistics import sys import time @@ -25,33 +23,28 @@ ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) - -def ensure_registries() -> None: - """Import env modules so they are registered in `unilab.envs.registry`.""" - try: - import unilab.envs.locomotion - - package = unilab.envs.locomotion - if hasattr(package, "__path__"): - for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): - try: - importlib.import_module(name) - except Exception: - pass - except ImportError: - pass - - -ensure_registries() - from unilab.algos.mlx.common import EmpiricalDiscountedVariationNormalization, RolloutBuffer from unilab.algos.mlx.ppo import MLPActorCritic, PPOConfig, PPOTrainer -from unilab.base import registry -from unilab.utils import render_many +from unilab.training import ( + create_env, + ensure_registries, + get_log_root, + parse_checkpoint_path, + render_play_mode, + setup_logger, +) +from unilab.training import ( + get_latest_checkpoint as get_latest_checkpoint_common, +) +from unilab.training import ( + get_latest_run as get_latest_run_common, +) from unilab.utils.experiment_tracking import ExperimentTracker from unilab.utils.obs_utils import flatten_obs_dict from unilab.utils.onpolicy_logger import OnPolicyLogger +ensure_registries() + TASK_STEP_TUNING = { "Go1JoystickFlatTerrain": {"threads": "32", "chunk": "4"}, "Go2JoystickFlatTerrain": {"threads": "56", "chunk": "16"}, @@ -84,22 +77,11 @@ def close(self) -> None: def get_latest_run(log_dir: Path) -> Path | None: - """Find latest run directory under a task log root.""" - if not log_dir.exists(): - return None - runs = sorted([p for p in log_dir.iterdir() if p.is_dir()]) - return runs[-1] if runs else None + return get_latest_run_common(log_dir) def get_latest_checkpoint(run_dir: Path) -> Path | None: - """Find the latest model_*.safetensors checkpoint in a run dir.""" - if not run_dir.exists(): - return None - model_files = [p for p in run_dir.glob("model_*.safetensors") if p.is_file()] - if not model_files: - return None - model_files.sort(key=lambda p: int(p.stem.split("_")[1])) - return model_files[-1] + return get_latest_checkpoint_common(run_dir, suffix=".safetensors") def save_trainer_state(path: Path, trainer: PPOTrainer, iteration: int) -> None: @@ -154,13 +136,7 @@ def _get_physics_state_snapshot(env) -> np.ndarray: def _get_log_root(cfg: DictConfig) -> Path: - """Get log root directory from algo_log_name config.""" - log_root = Path(cfg.training.log_root) if cfg.training.log_root else None - if log_root is None: - log_root = ROOT_DIR / "logs" / cfg.algo.algo_log_name - elif not log_root.is_absolute(): - log_root = ROOT_DIR / log_root - return log_root + return get_log_root(ROOT_DIR, cfg) def play_mlx_ppo(cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: str) -> str | None: @@ -171,10 +147,9 @@ def play_mlx_ppo(cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: s play_model_dtype = mx.float32 if use_fp16 else dtype play_env_num = cfg.training.play_env_num - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=play_env_num, - sim_backend=resolved_sim_backend, env_cfg_override=env_cfg_override, ) obs_dim = sum(env.obs_groups_spec.values()) @@ -182,31 +157,9 @@ def play_mlx_ppo(cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: s model = build_model(cfg.algo, obs_dim, action_dim, dtype=play_model_dtype) task_log_root = _get_log_root(cfg) / cfg.training.task_name - load_path: Path | None = None - load_run = cfg.algo.load_run - if load_run == "-1": - latest_run = get_latest_run(task_log_root) - if latest_run is not None: - load_path = get_latest_checkpoint(latest_run) - run_dir = latest_run - else: - run_dir = None - else: - candidate = Path(load_run) - if not candidate.exists(): - candidate = task_log_root / load_run - if candidate.is_dir(): - load_path = get_latest_checkpoint(candidate) - run_dir = candidate - elif candidate.is_file(): - load_path = candidate - run_dir = candidate.parent - else: - load_path = None - run_dir = None - - if load_path is None or not load_path.exists(): - print(f"Could not find valid model checkpoint from load_run={load_run}") + load_path, run_dir = parse_checkpoint_path(cfg, root_dir=ROOT_DIR, suffix=".safetensors") + if load_path is None or run_dir is None or not load_path.exists(): + print(f"Could not find valid model checkpoint from load_run={cfg.algo.load_run}") env.close() return None @@ -219,38 +172,31 @@ def play_mlx_ppo(cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: s obs_dict_play, _ = env.reset(play_reset_indices) obs = mx.array(flatten_obs_dict(obs_dict_play)) + def _play_step(current_obs): + obs_for_model = ( + current_obs.astype(play_model_dtype) + if getattr(current_obs, "dtype", None) != play_model_dtype + else current_obs + ) + actions_mx = model.policy(obs_for_model) + actions = mx.where(mx.isfinite(actions_mx), actions_mx, mx.zeros_like(actions_mx)) + actions = actions.astype(dtype) if getattr(actions, "dtype", None) != dtype else actions + state = env.step(np.asarray(actions)) + raw_obs = flatten_obs_dict(state.obs) + return mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0) + if resolved_sim_backend == "motrix": print("[MLX PPO] Starting interactive visualization (motrix native renderer)...") print("[MLX PPO] Close the render window to exit.") - env._backend.init_renderer() - - last_render_time = time.perf_counter() - render_dt = 1.0 / 60.0 try: - while True: - obs_for_model = ( - obs.astype(play_model_dtype) - if getattr(obs, "dtype", None) != play_model_dtype - else obs - ) - actions_mx = model.policy(obs_for_model) - actions = mx.where(mx.isfinite(actions_mx), actions_mx, mx.zeros_like(actions_mx)) - actions = ( - actions.astype(dtype) if getattr(actions, "dtype", None) != dtype else actions - ) - env_actions = np.asarray(actions) - state = env.step(env_actions) - raw_obs = flatten_obs_dict(state.obs) - obs = mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0) - - current_time = time.perf_counter() - elapsed = current_time - last_render_time - if elapsed < render_dt: - time.sleep(render_dt - elapsed) - last_render_time = time.perf_counter() - - env._backend.render() + render_play_mode( + env, + sim_backend="motrix", + num_steps=None, + initialize=lambda: obs, + step=_play_step, + ) except Exception as e: if "RenderClosedError" in str(type(e).__name__): print("[MLX PPO] Render window closed.") @@ -262,36 +208,21 @@ def play_mlx_ppo(cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: s output_dir = run_dir if run_dir is not None else task_log_root output_video = output_dir / "play_video.mp4" - state_list = [] print("[MLX PPO] Collecting physics states for play...") - for _ in range(cfg.training.play_steps): - obs_for_model = ( - obs.astype(play_model_dtype) if getattr(obs, "dtype", None) != play_model_dtype else obs - ) - actions_mx = model.policy(obs_for_model) - actions = mx.where(mx.isfinite(actions_mx), actions_mx, mx.zeros_like(actions_mx)) - actions = actions.astype(dtype) if getattr(actions, "dtype", None) != dtype else actions - env_actions = np.asarray(actions) - state = env.step(env_actions) - raw_obs = flatten_obs_dict(state.obs) - obs = mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0) - state_list.append(_get_physics_state_snapshot(env)) - - print(f"[MLX PPO] Rendering video to {output_video} ...") - frames = render_many.render_states_get_frames( - state_list, - env.cfg.model_file, - width=1280, - height=720, - camera_id=-1, - ) try: - import mediapy as media + render_play_mode( + env, + sim_backend=resolved_sim_backend, + num_steps=cfg.training.play_steps, + output_video=output_video, + initialize=lambda: obs, + step=_play_step, + frame_state_getter=lambda: _get_physics_state_snapshot(env), + ) except ImportError: print("mediapy is required for play video export. Install with `pip install mediapy`.") env.close() return None - media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt)) print(f"[MLX PPO] Play video saved: {output_video}") env.close() return str(output_video) @@ -328,15 +259,10 @@ def main(cfg: DictConfig) -> None: # TRAIN MODE log_dir = task_log_root / f"{timestamp}_{resolved_sim_backend}" - log_dir.mkdir(parents=True, exist_ok=True) - log_file_path = log_dir / "train.log" - log_fp = log_file_path.open("a", encoding="utf-8") + script_logger = setup_logger(log_dir, "mlx_ppo", echo=cfg.training.logger != "no_print") def log(msg: str) -> None: - if cfg.training.logger != "no_print": - print(msg) - log_fp.write(msg + "\n") - log_fp.flush() + script_logger.info(msg) tracker = ExperimentTracker( root_dir=ROOT_DIR, @@ -354,10 +280,9 @@ def log(msg: str) -> None: env_cfg_override = extract_reward_config(cfg) - env = registry.make( - task_name, + env = create_env( + cfg, num_envs=cfg.algo.num_envs, - sim_backend=resolved_sim_backend, env_cfg_override=env_cfg_override, ) if env.state is None: @@ -414,15 +339,7 @@ def log(msg: str) -> None: load_run = cfg.algo.load_run if load_run != "-1": - resume_candidate = Path(load_run) - if not resume_candidate.exists(): - resume_candidate = task_log_root / load_run - if resume_candidate.is_dir(): - ckpt = get_latest_checkpoint(resume_candidate) - elif resume_candidate.is_file(): - ckpt = resume_candidate - else: - ckpt = None + ckpt, _ = parse_checkpoint_path(cfg, root_dir=ROOT_DIR, suffix=".safetensors") if ckpt is not None and ckpt.exists(): model.load_weights(str(ckpt), strict=True) log(f"[MLX PPO] resumed_from={ckpt}") @@ -680,7 +597,6 @@ def log(msg: str) -> None: tracker.log_video(play_video_path) tracker.finish() - log_fp.close() if __name__ == "__main__": diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 97f2ee013..7784c1311 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -8,12 +8,21 @@ from pathlib import Path import hydra -from omegaconf import DictConfig, OmegaConf +from omegaconf import DictConfig ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) -from unilab.utils.algo_utils import ensure_registries +from unilab.training import ( + BackendAdapter, + create_env, + ensure_registries, + get_log_root, + render_play_mode, +) +from unilab.training import ( + resolve_checkpoint_path as resolve_checkpoint_path_common, +) from unilab.utils.experiment_tracking import ExperimentTracker @@ -31,52 +40,15 @@ def default_device(torch_module, preferred: str | None = None) -> str: def resolve_checkpoint_path( root_dir: Path, algo_log_name: str, task: str, load_run: str ) -> tuple[str | None, str | None]: - """Resolve latest or explicit checkpoint path for play mode.""" - base_log_dir = os.path.join(root_dir, "logs", algo_log_name, task) - - load_path = None - load_path_dir = None - if load_run == "-1": - if os.path.exists(base_log_dir): - all_runs = sorted( - [ - d - for d in os.listdir(base_log_dir) - if os.path.isdir(os.path.join(base_log_dir, d)) - ] - ) - if all_runs: - latest_run_dir = os.path.join(base_log_dir, all_runs[-1]) - model_files = sorted( - [ - f - for f in os.listdir(latest_run_dir) - if f.startswith("model_") and f.endswith(".pt") - ], - key=lambda x: int(x.split("_")[1].split(".")[0]), - ) - if model_files: - load_path = os.path.join(latest_run_dir, model_files[-1]) - load_path_dir = latest_run_dir - elif os.path.exists(load_run): - load_path = load_run - load_path_dir = os.path.dirname(load_path) - else: - potential_dir = os.path.join(base_log_dir, load_run) - if os.path.isdir(potential_dir): - model_files = sorted( - [ - f - for f in os.listdir(potential_dir) - if f.startswith("model_") and f.endswith(".pt") - ], - key=lambda x: int(x.split("_")[1].split(".")[0]), - ) - if model_files: - load_path = os.path.join(potential_dir, model_files[-1]) - load_path_dir = potential_dir - - return load_path, load_path_dir + checkpoint_path, checkpoint_dir = resolve_checkpoint_path_common( + Path(root_dir) / "logs" / algo_log_name / task, + load_run, + suffix=".pt", + ) + return ( + str(checkpoint_path) if checkpoint_path is not None else None, + str(checkpoint_dir) if checkpoint_dir is not None else None, + ) def extract_reset_obs(reset_result): @@ -102,87 +74,8 @@ def extract_play_obs(obs_dict): return obs_out -def _explicit_cli_override_keys() -> set[str]: - return {arg.split("=", 1)[0] for arg in sys.argv[1:] if "=" in arg} - - -def _apply_task_algo_overrides( - target, overrides, *, base_path: str, explicit_keys: set[str] -) -> None: - if OmegaConf.is_config(overrides): - override_dict = OmegaConf.to_container(overrides, resolve=True) - elif isinstance(overrides, dict): - override_dict = overrides - else: - return - - for key, value in override_dict.items(): - path = f"{base_path}.{key}" - if isinstance(value, dict): - if path in explicit_keys: - continue - if key not in target or target[key] is None: - target[key] = {} - _apply_task_algo_overrides( - target[key], value, base_path=path, explicit_keys=explicit_keys - ) - continue - if path not in explicit_keys: - target[key] = value - - -def _apply_task_env_profile( - env_cfg_override: dict, env_profile, *, explicit_keys: set[str] -) -> None: - if OmegaConf.is_config(env_profile): - env_profile_dict = OmegaConf.to_container(env_profile, resolve=True) - elif isinstance(env_profile, dict): - env_profile_dict = env_profile - else: - return - - reward_explicit = "reward" in explicit_keys or any( - k.startswith("reward.") for k in explicit_keys - ) - for key, value in env_profile_dict.items(): - if key == "reward_config" and reward_explicit: - continue - env_cfg_override[key] = value - - def build_task_motrix_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict | None: - from unilab.utils.reward_utils import extract_reward_config - - env_cfg_override = extract_reward_config(cfg) - explicit_keys = _explicit_cli_override_keys() - motrix_legacy = getattr(cfg, "motrix_legacy", None) - if motrix_legacy is None or not motrix_legacy.enabled or cfg.training.sim_backend != "motrix": - return env_cfg_override - - applies_to = getattr(motrix_legacy, "applies_to", None) - if applies_to is not None: - target_algo = getattr(applies_to, "algo", None) - if target_algo is not None and target_algo != algo_name: - return env_cfg_override - - algo_overrides = getattr(motrix_legacy, "algo_overrides", None) - if algo_overrides is not None: - _apply_task_algo_overrides( - cfg.algo, - algo_overrides, - base_path="algo", - explicit_keys=explicit_keys, - ) - - env_profile = getattr(motrix_legacy, "env_cfg_override", None) - if env_profile is not None: - _apply_task_env_profile( - env_cfg_override, - env_profile, - explicit_keys=explicit_keys, - ) - - return env_cfg_override + return BackendAdapter(cfg, root_dir=ROOT_DIR, algo_name=algo_name).build_task_env_cfg_override() def resolve_sac_use_symmetry(cfg: DictConfig) -> bool: @@ -202,15 +95,12 @@ def build_runner(algo_name: str, cfg: DictConfig): # Multi-GPU path if cfg.training.num_gpus > 1: from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner - from unilab.base import registry - from unilab.utils.algo_utils import ensure_registries ensure_registries() device = cfg.training.device or get_default_device() - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=1, - sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) assert env.action_space.shape @@ -334,12 +224,9 @@ def build_runner(algo_name: str, cfg: DictConfig): def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: """Play pipeline for off-policy algorithms.""" - import mediapy as media import numpy as np import torch - from unilab.base import registry - from unilab.utils import render_many from unilab.utils.algo_utils import build_actor env_cfg_override = build_task_motrix_offpolicy_env_cfg_override(algo_name, cfg) @@ -347,10 +234,9 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: device = default_device(torch, cfg.training.device) print(f"Using device for play: {device}") - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=cfg.training.play_env_num, - sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) obs_dim = resolve_play_obs_dim(env.obs_groups_spec) @@ -390,7 +276,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: ROOT_DIR, cfg.algo.algo_log_name, cfg.training.task_name, - cfg.training.load_run, + cfg.algo.load_run, ) if not load_path or not os.path.exists(load_path): print(f"Could not find checkpoint. load_path={load_path}") @@ -407,46 +293,42 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: normalizer.load_state_dict(checkpoint["obs_normalizer"]) normalizer.eval() - output_video = os.path.join(load_path_dir, "play_video.mp4") - if env.state is None: env.init_state() - env_indices = np.arange(cfg.training.play_env_num, dtype=np.int32) - obs_out = extract_reset_obs(env.reset(env_indices)) - obs_np = np.asarray(extract_play_obs(obs_out), dtype=np.float32) + + def _policy_step(obs_np: np.ndarray) -> np.ndarray: + obs_torch = torch.from_numpy(obs_np).to(device) + if normalizer: + obs_torch = normalizer(obs_torch, update=False) + actions_np = ( + actor.explore(obs_torch, deterministic=True).cpu().numpy() + if algo_name == "sac" + else actor(obs_torch).cpu().numpy() + ) + state = env.step(actions_np) + return np.asarray(extract_play_obs(state.obs), dtype=np.float32) # Use Motrix native rendering if cfg.training.sim_backend == "motrix": print("Starting interactive visualization (motrix native renderer)...") print("Close the render window to exit.") - env._backend.init_renderer() - - import time - - last_render_time = time.perf_counter() - render_dt = 1.0 / 60.0 with torch.inference_mode(): try: - while True: - obs_torch = torch.from_numpy(obs_np).to(device) - if normalizer: - obs_torch = normalizer(obs_torch, update=False) - actions_np = ( - actor.explore(obs_torch, deterministic=True).cpu().numpy() - if algo_name == "sac" - else actor(obs_torch).cpu().numpy() - ) - state = env.step(actions_np) - obs_np = np.asarray(extract_play_obs(state.obs), dtype=np.float32) - - current_time = time.perf_counter() - elapsed = current_time - last_render_time - if elapsed < render_dt: - time.sleep(render_dt - elapsed) - last_render_time = time.perf_counter() - - env._backend.render() + render_play_mode( + env, + sim_backend="motrix", + num_steps=None, + initialize=lambda: np.asarray( + extract_play_obs( + extract_reset_obs( + env.reset(np.arange(cfg.training.play_env_num, dtype=np.int32)) + ) + ), + dtype=np.float32, + ), + step=_policy_step, + ) except Exception as e: if "RenderClosedError" in str(type(e).__name__): print("Render window closed.") @@ -454,31 +336,26 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: raise return None - # MuJoCo backend: render to video output_video = os.path.join(load_path_dir, "play_video.mp4") - state_list = [] - print("Collecting physics states...") with torch.inference_mode(): - for _ in range(cfg.training.play_steps): - obs_torch = torch.from_numpy(obs_np).to(device) - if normalizer: - obs_torch = normalizer(obs_torch, update=False) - actions_np = ( - actor.explore(obs_torch, deterministic=True).cpu().numpy() - if algo_name == "sac" - else actor(obs_torch).cpu().numpy() - ) - state = env.step(actions_np) - obs_np = np.asarray(extract_play_obs(state.obs), dtype=np.float32) - state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy()) - - print("Rendering frames...") - frames = render_many.render_states_get_frames( - state_list, env.cfg.model_file, width=1280, height=720, camera_id=-1 - ) + render_play_mode( + env, + sim_backend=cfg.training.sim_backend, + num_steps=cfg.training.play_steps, + output_video=output_video, + initialize=lambda: np.asarray( + extract_play_obs( + extract_reset_obs( + env.reset(np.arange(cfg.training.play_env_num, dtype=np.int32)) + ) + ), + dtype=np.float32, + ), + step=_policy_step, + frame_state_getter=lambda: env._backend.get_physics_state(), + ) print(f"Saving video to {output_video} ...") - media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt)) print("Done.") return output_video @@ -492,12 +369,8 @@ def main(cfg: DictConfig) -> None: if cfg.training.log_dir is None: timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - log_dir = os.path.join( - ROOT_DIR, - "logs", - cfg.algo.algo_log_name, - task_name, - f"{timestamp}_{cfg.training.sim_backend}", + log_dir = str( + get_log_root(ROOT_DIR, cfg) / task_name / f"{timestamp}_{cfg.training.sim_backend}" ) else: log_dir = cfg.training.log_dir diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index ee053747f..781721574 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -1,17 +1,12 @@ import datetime -import importlib -import os -import pkgutil import statistics import sys import time from pathlib import Path import hydra -import numpy as np import torch from omegaconf import DictConfig, OmegaConf -from tensordict import TensorDict ROOT_DIR = Path(__file__).parent.parent SRC_DIR = ROOT_DIR / "src" @@ -20,12 +15,16 @@ if str(ROOT_DIR) not in sys.path: sys.path.insert(0, str(ROOT_DIR)) - +from unilab.training import ( + BackendAdapter, + create_env, + ensure_registries, + get_log_root, + parse_checkpoint_path, + render_play_mode, +) from unilab.utils.experiment_tracking import ExperimentTracker, patch_rsl_rl_wandb_writer -from unilab.utils.obs_utils import flatten_obs_dict -from unilab.utils.reward_utils import resolve_reward_dict from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper -from unilab.utils.torch_utils import to_numpy, to_torch from unilab.utils.xml_utils import materialize_scene_visual_override try: @@ -34,164 +33,24 @@ print("Could not import rsl_rl. Please ensure it is installed.") sys.exit(1) -from unilab.utils.rsl_rl_compat import ( - convert_config_v3_to_v4, - convert_config_v5, - is_rsl_rl_v4, - is_rsl_rl_v5, -) -from unilab.utils.run_utils import get_latest_run - - -def ensure_registries(): - for pkg_name in ( - "unilab.envs.locomotion", - "unilab.envs.manipulation", - "unilab.envs.motion_tracking", - ): - try: - package = importlib.import_module(pkg_name) - if hasattr(package, "__path__"): - for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): - try: - importlib.import_module(name) - except Exception: - pass - except ImportError: - pass - - -def _explicit_cli_override_keys() -> set[str]: - return {arg.split("=", 1)[0] for arg in sys.argv[1:] if "=" in arg} +from unilab.utils.rsl_rl_compat import convert_config_v5, is_rsl_rl_v5 -def _apply_task_algo_overrides( - target, overrides, *, base_path: str, explicit_keys: set[str] -) -> None: - if OmegaConf.is_config(overrides): - override_dict = OmegaConf.to_container(overrides, resolve=True) - elif isinstance(overrides, dict): - override_dict = overrides - else: - return - - for key, value in override_dict.items(): - path = f"{base_path}.{key}" - if isinstance(value, dict): - if path in explicit_keys: - continue - if key not in target or target[key] is None: - target[key] = {} - _apply_task_algo_overrides( - target[key], value, base_path=path, explicit_keys=explicit_keys - ) - continue - if path not in explicit_keys: - target[key] = value - - -def _apply_task_env_profile( - env_cfg_override: dict, env_profile, *, explicit_keys: set[str] -) -> None: - if OmegaConf.is_config(env_profile): - env_profile_dict = OmegaConf.to_container(env_profile, resolve=True) - elif isinstance(env_profile, dict): - env_profile_dict = env_profile - else: - return - - reward_explicit = "reward" in explicit_keys or any( - k.startswith("reward.") for k in explicit_keys +def _backend_adapter(cfg: DictConfig) -> BackendAdapter: + return BackendAdapter( + cfg, + root_dir=ROOT_DIR, + algo_name="ppo", + scene_materializer=materialize_scene_visual_override, ) - for key, value in env_profile_dict.items(): - if key == "reward_config" and reward_explicit: - continue - env_cfg_override[key] = value def build_task_motrix_ppo_env_cfg_override(cfg: DictConfig) -> dict: - env_cfg_override = {} - motrix_legacy = getattr(cfg, "motrix_legacy", None) - explicit_keys = _explicit_cli_override_keys() - env_cfg_override["reward_config"] = resolve_reward_dict(cfg) - if motrix_legacy is None or not motrix_legacy.enabled or cfg.training.sim_backend != "motrix": - return env_cfg_override - - algo_overrides = getattr(motrix_legacy, "algo_overrides", None) - if algo_overrides is not None: - _apply_task_algo_overrides( - cfg.algo, - algo_overrides, - base_path="algo", - explicit_keys=explicit_keys, - ) - - env_profile = getattr(motrix_legacy, "env_cfg_override", None) - if env_profile is not None: - _apply_task_env_profile( - env_cfg_override, - env_profile, - explicit_keys=explicit_keys, - ) - return env_cfg_override - - -def _resolve_root_relative_path(path_value: str) -> str: - candidate = Path(path_value) - if candidate.is_absolute(): - return str(candidate) - return str((ROOT_DIR / candidate).resolve()) + return _backend_adapter(cfg).build_task_env_cfg_override() def build_motrix_play_ppo_env_cfg_override(cfg: DictConfig) -> dict: - env_cfg_override = build_task_motrix_ppo_env_cfg_override(cfg) - explicit_keys = _explicit_cli_override_keys() - play_profile = getattr(cfg, "motrix_play_only", None) - if ( - play_profile is None - or not getattr(play_profile, "enabled", False) - or cfg.training.sim_backend != "motrix" - or not cfg.training.play_only - ): - return env_cfg_override - - training_overrides = getattr(play_profile, "training_overrides", None) - if training_overrides is not None: - _apply_task_algo_overrides( - cfg.training, - training_overrides, - base_path="training", - explicit_keys=explicit_keys, - ) - - env_profile = getattr(play_profile, "env_cfg_override", None) - if env_profile is not None: - _apply_task_env_profile( - env_cfg_override, - env_profile, - explicit_keys=explicit_keys, - ) - - scene_override = getattr(play_profile, "scene_override", None) - if scene_override is None or not getattr(scene_override, "enabled", False): - return env_cfg_override - - source_model_file = getattr(scene_override, "source_model_file", None) - if not source_model_file: - raise ValueError("motrix_play_only.scene_override.source_model_file must be configured") - - env_cfg_override["model_file"] = materialize_scene_visual_override( - _resolve_root_relative_path(str(source_model_file)), - ground_texture_file=( - _resolve_root_relative_path(str(scene_override.ground_texture_file)) - if getattr(scene_override, "ground_texture_file", None) - else None - ), - ground_texrepeat=getattr(scene_override, "ground_texrepeat", None), - skybox_rgb1=getattr(scene_override, "skybox_rgb1", None), - skybox_rgb2=getattr(scene_override, "skybox_rgb2", None), - ) - return env_cfg_override + return _backend_adapter(cfg).build_play_env_cfg_override() def run_motrix_rsl_play_loop( @@ -201,47 +60,31 @@ def run_motrix_rsl_play_loop( render_spacing: float, num_steps: int | None = None, ) -> None: - import time - env = wrapped_env.env - env._backend.init_renderer(spacing=render_spacing) - obs, _ = wrapped_env.reset() - - last_render_time = time.perf_counter() - render_dt = 1.0 / 60.0 - steps_run = 0 with torch.inference_mode(): - while num_steps is None or steps_run < num_steps: - actions = policy(obs) - obs, _, _, _ = wrapped_env.step(actions) - - current_time = time.perf_counter() - elapsed = current_time - last_render_time - if elapsed < render_dt: - time.sleep(render_dt - elapsed) - last_render_time = time.perf_counter() - - env._backend.render() - steps_run += 1 + render_play_mode( + env, + sim_backend="motrix", + render_spacing=render_spacing, + num_steps=num_steps, + initialize=lambda: wrapped_env.reset()[0], + step=lambda obs: wrapped_env.step(policy(obs))[0], + ) def _get_log_root(cfg: DictConfig) -> str: - """Get log root directory from algo_log_name config.""" - return str(ROOT_DIR / "logs" / cfg.algo.algo_log_name) + return str(get_log_root(ROOT_DIR, cfg)) def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: """Play mode for RSL-RL.""" - from unilab.base import registry - env_cfg_override = build_motrix_play_ppo_env_cfg_override(cfg) - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=cfg.training.play_env_num, - sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) wrapped_env = RslRlVecEnvWrapper(env, device=device) @@ -252,38 +95,12 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: if is_rsl_rl_v5(): train_cfg = convert_config_v5(train_cfg) - log_root = _get_log_root(cfg) - base_log_dir = Path(log_root) / cfg.training.task_name - load_path = None - - load_run = cfg.algo.load_run - if load_run == "-1": - load_path = get_latest_run(str(base_log_dir)) - else: - if os.path.exists(load_run): - load_path = load_run - else: - load_path = str(base_log_dir / load_run) - - if not load_path or not os.path.exists(load_path): + load_path, load_path_dir = parse_checkpoint_path(cfg, root_dir=ROOT_DIR) + if load_path is None or load_path_dir is None or not load_path.exists(): print(f"Could not find run to load at {load_path}") return None - if os.path.isdir(load_path): - model_files = [ - f for f in os.listdir(load_path) if f.startswith("model_") and f.endswith(".pt") - ] - if len(model_files) > 0: - model_files.sort(key=lambda x: int(x.split("_")[1].split(".")[0])) - load_path_dir = load_path - load_path = os.path.join(load_path, model_files[-1]) - print(f"Loading latest model: {load_path}") - else: - print(f"No model files found in {load_path}") - return None - else: - load_path_dir = os.path.dirname(load_path) - + print(f"Loading latest model: {load_path}") _ckpt_keys = set(torch.load(load_path, map_location="cpu", weights_only=True).keys()) if "actor_state_dict" not in _ckpt_keys: print( @@ -293,7 +110,7 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: return None runner = OnPolicyRunner(wrapped_env, train_cfg, log_dir=None, device=device) - runner.load(load_path) + runner.load(str(load_path)) policy = runner.get_inference_policy(device=device) if cfg.training.sim_backend == "motrix": @@ -312,37 +129,25 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: else: raise else: - import mediapy as media - - from unilab.utils import render_many - output_video = Path(load_path_dir) / "play_video.mp4" print(f"Rendering video to {output_video}...") - obs, _ = wrapped_env.reset() - state_list = [] - print("Collecting physics states...") with torch.inference_mode(): - for _ in range(cfg.training.play_steps): - actions = policy(obs) - obs, _, _, _ = wrapped_env.step(actions) - state_list.append(to_numpy(env._backend.get_physics_state()).copy()) - - print("Rendering frames...") - frames = render_many.render_states_get_frames( - state_list, - env.cfg.model_file, - width=1280, - height=720, - camera_id=-1, - cam_distance=cfg.training.cam_distance, - cam_elevation=cfg.training.cam_elevation, - cam_azimuth=cfg.training.cam_azimuth, - ) - - print(f"Saving video to {output_video} with mediapy...") - media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt)) + render_play_mode( + env, + sim_backend=cfg.training.sim_backend, + num_steps=cfg.training.play_steps, + output_video=output_video, + initialize=lambda: wrapped_env.reset()[0], + step=lambda obs: wrapped_env.step(policy(obs))[0], + frame_state_getter=lambda: env._backend.get_physics_state(), + camera_kwargs={ + "cam_distance": cfg.training.cam_distance, + "cam_elevation": cfg.training.cam_elevation, + "cam_azimuth": cfg.training.cam_azimuth, + }, + ) print("Done.") return str(output_video) @@ -353,10 +158,6 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: def main(cfg: DictConfig) -> None: ensure_registries() - from omegaconf import OmegaConf - - from unilab.base import registry - env_cfg_override = build_task_motrix_ppo_env_cfg_override(cfg) if torch.cuda.is_available(): @@ -402,10 +203,9 @@ def main(cfg: DictConfig) -> None: try: if not cfg.training.play_only: - env = registry.make( - cfg.training.task_name, + env = create_env( + cfg, num_envs=cfg.algo.num_envs, - sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) wrapped_env = RslRlVecEnvWrapper(env, device=device) @@ -437,17 +237,10 @@ def main(cfg: DictConfig) -> None: runner = OnPolicyRunner(wrapped_env, train_cfg, log_dir=log_dir, device=device) if cfg.algo.load_run != "-1": - if os.path.exists(cfg.algo.load_run): - resume_path = cfg.algo.load_run - else: - log_root = _get_log_root(cfg) - base_log_dir = Path(log_root) / cfg.training.task_name - run_path = base_log_dir / cfg.algo.load_run - resume_path = str(run_path) if run_path.exists() else None - + resume_path, _ = parse_checkpoint_path(cfg, root_dir=ROOT_DIR) if resume_path: print(f"Resuming from {resume_path}") - runner.load(resume_path) + runner.load(str(resume_path)) train_start_wall = time.time() runner.learn(num_learning_iterations=max_iterations, init_at_random_ep_len=True) diff --git a/src/unilab/training/__init__.py b/src/unilab/training/__init__.py new file mode 100644 index 000000000..0c9f6de11 --- /dev/null +++ b/src/unilab/training/__init__.py @@ -0,0 +1,27 @@ +"""Shared training helpers for entrypoint scripts.""" + +from unilab.training.backend_adapter import BackendAdapter +from unilab.training.common import ( + create_env, + ensure_registries, + get_latest_checkpoint, + get_latest_run, + get_log_root, + parse_checkpoint_path, + render_play_mode, + resolve_checkpoint_path, + setup_logger, +) + +__all__ = [ + "BackendAdapter", + "create_env", + "ensure_registries", + "get_latest_checkpoint", + "get_latest_run", + "get_log_root", + "parse_checkpoint_path", + "render_play_mode", + "resolve_checkpoint_path", + "setup_logger", +] diff --git a/src/unilab/training/backend_adapter.py b/src/unilab/training/backend_adapter.py new file mode 100644 index 000000000..0cfd1948d --- /dev/null +++ b/src/unilab/training/backend_adapter.py @@ -0,0 +1,161 @@ +"""Backend-specific config adaptation for training entrypoints.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, Callable + +from omegaconf import DictConfig, OmegaConf + +from unilab.utils.reward_utils import extract_reward_config +from unilab.utils.xml_utils import materialize_scene_visual_override + + +class BackendAdapter: + """Apply backend-specific config and env overrides while respecting CLI overrides.""" + + def __init__( + self, + cfg: DictConfig, + *, + root_dir: str | Path, + algo_name: str | None = None, + scene_materializer: Callable[..., str] = materialize_scene_visual_override, + ) -> None: + self.cfg = cfg + self.root_dir = Path(root_dir) + self.algo_name = algo_name + self.scene_materializer = scene_materializer + self.explicit_keys = { + arg.split("=", 1)[0].lstrip("+") for arg in sys.argv[1:] if "=" in arg + } + + def build_task_env_cfg_override(self) -> dict[str, Any]: + """Build env_cfg_override for training/play entrypoints.""" + env_cfg_override = extract_reward_config(self.cfg) + motrix_legacy = getattr(self.cfg, "motrix_legacy", None) + if ( + motrix_legacy is None + or not getattr(motrix_legacy, "enabled", False) + or self.cfg.training.sim_backend != "motrix" + ): + return env_cfg_override + + applies_to = getattr(motrix_legacy, "applies_to", None) + target_algo = getattr(applies_to, "algo", None) if applies_to is not None else None + if self.algo_name is not None and target_algo is not None and target_algo != self.algo_name: + return env_cfg_override + + algo_overrides = getattr(motrix_legacy, "algo_overrides", None) + if algo_overrides is not None: + self._apply_nested_overrides( + self.cfg.algo, + algo_overrides, + base_path="algo", + ) + + env_profile = getattr(motrix_legacy, "env_cfg_override", None) + if env_profile is not None: + self._apply_env_profile(env_cfg_override, env_profile) + + return env_cfg_override + + def build_play_env_cfg_override(self) -> dict[str, Any]: + """Build play-mode overrides, including Motrix scene customization when configured.""" + env_cfg_override = self.build_task_env_cfg_override() + play_profile = getattr(self.cfg, "motrix_play_only", None) + if ( + play_profile is None + or not getattr(play_profile, "enabled", False) + or self.cfg.training.sim_backend != "motrix" + or not self.cfg.training.play_only + ): + return env_cfg_override + + training_overrides = getattr(play_profile, "training_overrides", None) + if training_overrides is not None: + self._apply_nested_overrides( + self.cfg.training, + training_overrides, + base_path="training", + ) + + env_profile = getattr(play_profile, "env_cfg_override", None) + if env_profile is not None: + self._apply_env_profile(env_cfg_override, env_profile) + + scene_override = getattr(play_profile, "scene_override", None) + if scene_override is None or not getattr(scene_override, "enabled", False): + return env_cfg_override + + source_model_file = getattr(scene_override, "source_model_file", None) + if not source_model_file: + raise ValueError("motrix_play_only.scene_override.source_model_file must be configured") + + env_cfg_override["model_file"] = self.scene_materializer( + self._resolve_root_relative_path(str(source_model_file)), + ground_texture_file=( + self._resolve_root_relative_path(str(scene_override.ground_texture_file)) + if getattr(scene_override, "ground_texture_file", None) + else None + ), + ground_texrepeat=getattr(scene_override, "ground_texrepeat", None), + skybox_rgb1=getattr(scene_override, "skybox_rgb1", None), + skybox_rgb2=getattr(scene_override, "skybox_rgb2", None), + ) + return env_cfg_override + + def _apply_nested_overrides(self, target: Any, overrides: Any, *, base_path: str) -> None: + if OmegaConf.is_config(overrides): + override_dict = OmegaConf.to_container(overrides, resolve=True) + elif isinstance(overrides, dict): + override_dict = overrides + else: + return + + if not isinstance(override_dict, dict): + return + + for key, value in override_dict.items(): + key_str = str(key) + path = f"{base_path}.{key_str}" + if isinstance(value, dict): + if path in self.explicit_keys: + continue + if key_str not in target or target[key_str] is None: + target[key_str] = {} + self._apply_nested_overrides(target[key_str], value, base_path=path) + continue + if path not in self.explicit_keys: + target[key_str] = value + + def _apply_env_profile(self, env_cfg_override: dict[str, Any], env_profile: Any) -> None: + if OmegaConf.is_config(env_profile): + env_profile_dict = OmegaConf.to_container(env_profile, resolve=True) + elif isinstance(env_profile, dict): + env_profile_dict = env_profile + else: + return + + if not isinstance(env_profile_dict, dict): + return + + reward_explicit = any( + key == "reward" + or key.startswith("reward.") + or key.startswith("reward@") + or key.startswith("reward_") + for key in self.explicit_keys + ) + for key, value in env_profile_dict.items(): + key_str = str(key) + if key_str == "reward_config" and reward_explicit: + continue + env_cfg_override[key_str] = value + + def _resolve_root_relative_path(self, path_value: str) -> str: + candidate = Path(path_value) + if candidate.is_absolute(): + return str(candidate) + return str((self.root_dir / candidate).resolve()) diff --git a/src/unilab/training/common.py b/src/unilab/training/common.py new file mode 100644 index 000000000..9b475b79f --- /dev/null +++ b/src/unilab/training/common.py @@ -0,0 +1,225 @@ +"""Shared helpers for training entrypoints.""" + +from __future__ import annotations + +import logging +import time +from pathlib import Path +from typing import Any, Callable, TypeVar + +import numpy as np +from omegaconf import DictConfig, OmegaConf + +from unilab.utils.algo_utils import ensure_registries as _ensure_registries + +ObsT = TypeVar("ObsT") + + +def ensure_registries() -> None: + """Import env modules so registry-based entrypoints can instantiate tasks.""" + _ensure_registries() + + +def get_log_root(root_dir: str | Path, cfg: DictConfig) -> Path: + """Resolve the algorithm log root, honoring optional training.log_root overrides.""" + configured_root = OmegaConf.select(cfg, "training.log_root") + if configured_root: + log_root = Path(str(configured_root)) + return log_root if log_root.is_absolute() else Path(root_dir) / log_root + return Path(root_dir) / "logs" / str(OmegaConf.select(cfg, "algo.algo_log_name")) + + +def get_latest_run(log_dir: str | Path) -> Path | None: + """Return the lexicographically latest run directory under a task log root.""" + base_dir = Path(log_dir) + if not base_dir.exists(): + return None + runs = sorted(path for path in base_dir.iterdir() if path.is_dir()) + return runs[-1] if runs else None + + +def get_latest_checkpoint(run_dir: str | Path, *, suffix: str = ".pt") -> Path | None: + """Return the latest model checkpoint inside a run directory.""" + run_path = Path(run_dir) + if not run_path.exists(): + return None + + def _iteration(path: Path) -> int: + stem_parts = path.stem.split("_", 1) + if len(stem_parts) != 2: + return -1 + try: + return int(stem_parts[1]) + except ValueError: + return -1 + + model_files = [ + path + for path in run_path.iterdir() + if path.is_file() and path.name.startswith("model_") and path.suffix == suffix + ] + if not model_files: + return None + return max(model_files, key=_iteration) + + +def resolve_checkpoint_path( + base_log_dir: str | Path, + load_run: str, + *, + suffix: str = ".pt", +) -> tuple[Path | None, Path | None]: + """Resolve a latest or explicit checkpoint path from a task log root.""" + base_dir = Path(base_log_dir) + if load_run == "-1": + run_dir = get_latest_run(base_dir) + if run_dir is None: + return None, None + checkpoint = get_latest_checkpoint(run_dir, suffix=suffix) + return (checkpoint, run_dir) if checkpoint is not None else (None, None) + + candidate = Path(load_run) + if not candidate.exists(): + candidate = base_dir / load_run + if candidate.is_file(): + return candidate, candidate.parent + if candidate.is_dir(): + checkpoint = get_latest_checkpoint(candidate, suffix=suffix) + return (checkpoint, candidate) if checkpoint is not None else (None, None) + return None, None + + +def parse_checkpoint_path( + cfg: DictConfig, + *, + root_dir: str | Path, + load_run: str | None = None, + task_name: str | None = None, + suffix: str = ".pt", +) -> tuple[Path | None, Path | None]: + """Resolve a checkpoint path from Hydra config and repository root.""" + selected_task = task_name or str(OmegaConf.select(cfg, "training.task_name")) + selected_run = load_run or str(OmegaConf.select(cfg, "algo.load_run", default="-1")) + base_log_dir = get_log_root(root_dir, cfg) / selected_task + return resolve_checkpoint_path(base_log_dir, selected_run, suffix=suffix) + + +def setup_logger( + log_dir: str | Path, + algo_name: str, + *, + echo: bool = True, + filename: str = "train.log", +) -> logging.Logger: + """Create a simple file-backed logger for script-local progress messages.""" + path = Path(log_dir) + path.mkdir(parents=True, exist_ok=True) + + logger_name = f"unilab.training.{algo_name}.{path.resolve()}" + logger = logging.getLogger(logger_name) + logger.setLevel(logging.INFO) + logger.propagate = False + + for handler in list(logger.handlers): + logger.removeHandler(handler) + handler.close() + + formatter = logging.Formatter("%(message)s") + + file_handler = logging.FileHandler(path / filename, encoding="utf-8") + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + if echo: + stream_handler = logging.StreamHandler() + stream_handler.setFormatter(formatter) + logger.addHandler(stream_handler) + + return logger + + +def create_env( + cfg: DictConfig, + *, + num_envs: int, + env_cfg_override: dict[str, Any] | None = None, + sim_backend: str | None = None, + task_name: str | None = None, +): + """Construct an environment via the registry using the current Hydra config.""" + from unilab.base import registry + + return registry.make( + task_name or str(OmegaConf.select(cfg, "training.task_name")), + num_envs=num_envs, + sim_backend=sim_backend or str(OmegaConf.select(cfg, "training.sim_backend")), + env_cfg_override=env_cfg_override, + ) + + +def render_play_mode( + env, + *, + sim_backend: str, + initialize: Callable[[], ObsT], + step: Callable[[ObsT], ObsT], + num_steps: int | None, + output_video: str | Path | None = None, + render_spacing: float | None = None, + frame_state_getter: Callable[[], np.ndarray] | None = None, + camera_kwargs: dict[str, Any] | None = None, +) -> str | None: + """Render interactive Motrix play or MuJoCo video generation through shared callbacks.""" + if sim_backend == "motrix": + if render_spacing is None: + env._backend.init_renderer() + else: + try: + env._backend.init_renderer(spacing=render_spacing) + except TypeError: + env._backend.init_renderer() + + obs = initialize() + last_render_time = time.perf_counter() + render_dt = 1.0 / 60.0 + steps_run = 0 + + while num_steps is None or steps_run < num_steps: + obs = step(obs) + current_time = time.perf_counter() + elapsed = current_time - last_render_time + if elapsed < render_dt: + time.sleep(render_dt - elapsed) + last_render_time = time.perf_counter() + env._backend.render() + steps_run += 1 + return None + + if num_steps is None: + raise ValueError("MuJoCo play rendering requires a finite num_steps value.") + if output_video is None: + raise ValueError("MuJoCo play rendering requires an output_video path.") + if frame_state_getter is None: + raise ValueError("MuJoCo play rendering requires a frame_state_getter callback.") + + obs = initialize() + state_list = [] + for _ in range(num_steps): + obs = step(obs) + state_list.append(np.asarray(frame_state_getter(), dtype=np.float32).copy()) + + from unilab.utils import render_many + + frames = render_many.render_states_get_frames( + state_list, + env.cfg.model_file, + width=1280, + height=720, + camera_id=-1, + **(camera_kwargs or {}), + ) + + import mediapy as media + + media.write_video(str(output_video), frames, fps=int(1.0 / env.cfg.ctrl_dt)) + return str(output_video) diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py new file mode 100644 index 000000000..cf1286e24 --- /dev/null +++ b/tests/training/test_training_helpers.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from unilab.training import ( + BackendAdapter, + get_latest_checkpoint, + get_latest_run, + parse_checkpoint_path, +) + +_ROOT_DIR = Path(__file__).resolve().parents[2] +_CONF_DIR = _ROOT_DIR / "conf" + + +def _ppo_cfg(overrides: list[str] | None = None): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): + return compose("config", overrides=overrides or []) + + +def _offpolicy_cfg(overrides: list[str] | None = None): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(_CONF_DIR / "offpolicy"), version_base="1.3"): + return compose("config", overrides=overrides or []) + + +def test_get_latest_run_and_checkpoint_support_shared_checkpoint_resolution(tmp_path: Path): + task_dir = tmp_path / "logs" / "custom_ppo" / "MyTask" + older_run = task_dir / "2024-01-01_00-00-00_mujoco" + newer_run = task_dir / "2024-02-01_00-00-00_mujoco" + older_run.mkdir(parents=True) + newer_run.mkdir(parents=True) + (older_run / "model_1.pt").write_bytes(b"") + (newer_run / "model_5.pt").write_bytes(b"") + (newer_run / "model_9.pt").write_bytes(b"") + + latest_run = get_latest_run(task_dir) + latest_checkpoint = get_latest_checkpoint(newer_run) + + assert latest_run == newer_run + assert latest_checkpoint == newer_run / "model_9.pt" + + +def test_parse_checkpoint_path_uses_algo_log_name_from_cfg(tmp_path: Path): + cfg = _ppo_cfg() + cfg.algo.algo_log_name = "custom_ppo" + + run_dir = ( + tmp_path / "logs" / "custom_ppo" / cfg.training.task_name / "2024-01-01_00-00-00_mujoco" + ) + run_dir.mkdir(parents=True) + model_path = run_dir / "model_12.pt" + model_path.write_bytes(b"") + + checkpoint_path, checkpoint_dir = parse_checkpoint_path(cfg, root_dir=tmp_path) + + assert checkpoint_path == model_path + assert checkpoint_dir == run_dir + + +def test_backend_adapter_applies_motrix_legacy_for_matching_algo(): + cfg = _offpolicy_cfg(["task=go1_joystick", "training.sim_backend=motrix"]) + + env_cfg_override = BackendAdapter( + cfg, root_dir=_ROOT_DIR, algo_name="sac" + ).build_task_env_cfg_override() + + assert cfg.algo.num_envs == 4096 + assert cfg.algo.max_iterations == 2000 + assert env_cfg_override["legacy_motrix_profile"]["enabled"] is True + assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) + + +def test_backend_adapter_builds_motrix_play_scene_override(): + cfg = _ppo_cfg( + ["task=g1_motion_tracking", "training.sim_backend=motrix", "training.play_only=true"] + ) + captured: dict[str, object] = {} + + def _fake_materializer(source_model_file: str, **kwargs) -> str: + captured["source_model_file"] = source_model_file + captured.update(kwargs) + return "/tmp/g1_motion_tracking_play_scene.xml" + + env_cfg_override = BackendAdapter( + cfg, + root_dir=_ROOT_DIR, + algo_name="ppo", + scene_materializer=_fake_materializer, + ).build_play_env_cfg_override() + + assert cfg.training.play_env_num == 128 + assert env_cfg_override["render_spacing"] == pytest.approx(2.5) + assert env_cfg_override["model_file"] == "/tmp/g1_motion_tracking_play_scene.xml" + assert captured["ground_texture_file"] == str( + _ROOT_DIR / "src/unilab/assets/robots/g1/floor.png" + )