# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import logging from pathlib import Path from absl import app, flags from skrl import config from motrix_rl import utils logger = logging.getLogger(__name__) _ENV = flags.DEFINE_string("env", "cartpole", "The env to play") _SIM_BACKEND = flags.DEFINE_string( "sim-backend", None, "The simulation backend to use.(If not specified, it will be choosen automatically)", ) _POLICY = flags.DEFINE_string("policy", None, "The policy to load") _NUM_ENVS = flags.DEFINE_integer("num-envs", 2048, "Number of envs to play") _SEED = flags.DEFINE_integer("seed", None, "Random seed for reproducibility") _RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed") _RLLIB = flags.DEFINE_string( "rllib", None, "The RL framework (skrl/rslrl). Auto-discovered from latest training if not specified." ) _FORCE_PHASE = flags.DEFINE_integer("force-phase", None, "Lock terrain phase (0=flat,1=rough,2=stairs,3=mixed)") def get_inference_backend(policy_path: Path | str, rllib: str): """Determine the backend from RL framework and policy file extension.""" if rllib == "rslrl": # RSLRL always uses torch backend return "torch" # Handle both Path and str types suffix = policy_path.suffix if isinstance(policy_path, Path) else Path(policy_path).suffix if suffix == ".pt": return "torch" if suffix == ".pickle": return "jax" else: raise Exception(f"Unknown policy format: {policy_path}") def discover_rllib(env_name: str) -> tuple[str, Path]: """ Discover the RL framework and best policy from the most recent training run. Args: env_name: The name of the environment Returns: Tuple of (RL framework name, path to best policy) Raises: FileNotFoundError: If no training results are found """ base_dir = Path(f"runs/{env_name}") if not base_dir.exists(): raise FileNotFoundError(f"No training results found for environment '{env_name}' in {base_dir}") frameworks = [] for framework in ["skrl", "rslrl"]: framework_dir = base_dir / framework if framework_dir.exists() and framework_dir.is_dir(): # Get all training run directories training_runs = [d for d in framework_dir.iterdir() if d.is_dir()] if training_runs: # Find the most recent run for this framework latest_run = max(training_runs, key=lambda x: x.stat().st_mtime) frameworks.append((framework, latest_run.stat().st_mtime, latest_run)) if not frameworks: raise FileNotFoundError(f"No training runs found for environment '{env_name}' in {base_dir}") # Return the framework with the most recent training run and its best policy latest_framework, _, latest_run_dir = max(frameworks, key=lambda x: x[1]) logger.info(f"Auto-discovered RL framework: {latest_framework}") # Find best policy in the latest run directory if latest_framework == "rslrl": # RSLRL uses model_*.pt format model_files = list(latest_run_dir.glob("model_*.pt")) if not model_files: raise FileNotFoundError(f"No policy files found in {latest_run_dir}") def extract_iteration(filename): stem = Path(filename).stem parts = stem.split("_") if len(parts) >= 2: try: return int(parts[1]) except ValueError: return 0 return 0 best_policy = max(model_files, key=lambda f: (f.stat().st_mtime, extract_iteration(f))) else: # SKRL uses checkpoints subdirectory checkpoints_dir = latest_run_dir / "checkpoints" if not checkpoints_dir.exists(): raise FileNotFoundError(f"No checkpoints directory found in {latest_run_dir}") # First, try to find best_agent files best_files = list(checkpoints_dir.glob("best_agent.*")) if best_files: best_policy = best_files[0] else: # Find checkpoint with highest timestep checkpoint_files = list(checkpoints_dir.glob("agent_*.pt")) + list(checkpoints_dir.glob("agent_*.pickle")) if not checkpoint_files: raise FileNotFoundError(f"No policy files found in {checkpoints_dir}") def extract_timestep(filename): stem = Path(filename).stem parts = stem.split("_") if len(parts) >= 2: try: return int(parts[1]) except ValueError: return 0 return 0 best_policy = max(checkpoint_files, key=extract_timestep) return latest_framework, best_policy def main(argv): device_supports = utils.get_device_supports() logger.info(device_supports) env_name = _ENV.value enable_render = True rl_override = {} if _NUM_ENVS.present: rl_override["play_num_envs"] = _NUM_ENVS.value if _RAND_SEED.value: rl_override["runner.seed"] = None elif _SEED.present: rl_override["runner.seed"] = _SEED.value sim_backend = _SIM_BACKEND.value rllib = None policy_path = None if _POLICY.present: if not _RLLIB.present: logger.error("Error: --policy specified but --rllib not specified") return rllib = _RLLIB.value policy_path = _POLICY.value logger.info(f"Using specified RL framework: {rllib}") logger.info(f"Using specified policy: {policy_path}") else: # if policy is not specified, search for the lastest training run and use its best policy try: rllib, policy_path = discover_rllib(env_name) logger.info(f"Auto-discovered RL framework: {rllib}") logger.info(f"Auto-discovered best policy: {policy_path}") except FileNotFoundError as e: logger.error(f"Error: {e}") logger.error("Please specify --rllib or train a model first") return backend = get_inference_backend(policy_path, rllib) # Build env config overrides env_cfg_override = {} if _FORCE_PHASE.present: env_cfg_override["force_phase"] = _FORCE_PHASE.value if not env_cfg_override: env_cfg_override = None if rllib == "rslrl": # RSLRL evaluation flow (always uses torch backend) assert device_supports.torch, "PyTorch is not available on your device" from motrix_rl.rslrl.torch.train import ppo config.torch.backend = "torch" trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render, env_cfg_override=env_cfg_override) trainer.play(policy_path) elif backend == "jax": assert device_supports.jax, "jax is not avaliable on your device " from motrix_rl.skrl.jax.train import ppo config.jax.backend = "jax" # or "numpy" trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render, env_cfg_override=env_cfg_override) trainer.play(policy_path) elif backend == "torch": assert device_supports.torch, "torch is not avaliable on your device" from motrix_rl.skrl.torch.train import ppo config.torch.backend = "torch" trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render, env_cfg_override=env_cfg_override) trainer.play(policy_path) if __name__ == "__main__": app.run(main)