221 lines
8.2 KiB
Python
221 lines
8.2 KiB
Python
# 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)
|