cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: release v0.3.0

Browse Source
This commit is contained in:
motphys-developers
2026-04-02 03:45:10 +00:00
parent c84d382b8c
commit e1421d1055
232 changed files with 20258 additions and 2004 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
scripts/play.py
View File

@@ -20,7 +20,6 @@ from absl import app, flags
from skrl import config
from motrix_rl import utils
from motrix_rl.skrl import get_log_dir
logger = logging.getLogger(__name__)
@@ -34,77 +33,109 @@ _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."
)
def get_inference_backend(policy_path: str):
if policy_path.endswith(".pt"):
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"
if policy_path.endswith(".pickle"):
# 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 find_best_policy(env_name: str) -> str:
def discover_rllib(env_name: str) -> tuple[str, Path]:
"""
Find the most recent best policy for the given environment.
Discover the RL framework and best policy from the most recent training run.
Args:
env_name: The name of the environment
Returns:
Path to the best policy file
Tuple of (RL framework name, path to best policy)
Raises:
FileNotFoundError: If no policy files are found
FileNotFoundError: If no training results are found
"""
# Base runs directory
base_dir = Path(f"runs/{env_name}")
env_dir = Path(get_log_dir(env_name))
if not base_dir.exists():
raise FileNotFoundError(f"No training results found for environment '{env_name}' in {base_dir}")
if not env_dir.exists():
raise FileNotFoundError(f"No training results found for environment '{env_name}' in {env_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))
# Find all training run directories (pattern: YY-MM-DD_HH-MM-SS-_XXXXX_PPO)
training_runs = [d for d in env_dir.iterdir() if d.is_dir()]
if not frameworks:
raise FileNotFoundError(f"No training runs found for environment '{env_name}' in {base_dir}")
if not training_runs:
raise FileNotFoundError(f"No training runs found for environment '{env_name}'")
# 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}")
# Sort by modification time to get the most recent
latest_run = max(training_runs, key=lambda x: x.stat().st_mtime)
checkpoints_dir = latest_run / "checkpoints"
# 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}")
if not checkpoints_dir.exists():
raise FileNotFoundError(f"No checkpoints directory found in {latest_run}")
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
# First, try to find best_agent files (highest performance models)
best_files = list(checkpoints_dir.glob("best_agent.*"))
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}")
if best_files:
# Return the first best_agent file found (there should only be one)
return str(best_files[0])
# 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}")
# If no best_agent files, find the checkpoint with the 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}")
# Extract timestep from filename and find the highest
def extract_timestep(filename):
# Pattern: agent_{timestep}.ext
stem = Path(filename).stem # agent_{timestep}
parts = stem.split("_")
if len(parts) >= 2:
try:
return int(parts[1])
except ValueError:
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
return 0
latest_checkpoint = max(checkpoint_files, key=extract_timestep)
return str(latest_checkpoint)
best_policy = max(checkpoint_files, key=extract_timestep)
return latest_framework, best_policy
def main(argv):
@@ -119,28 +150,45 @@ def main(argv):
rl_override["play_num_envs"] = _NUM_ENVS.value
if _RAND_SEED.value:
rl_override["seed"] = None
rl_override["runner.seed"] = None
elif _SEED.present:
rl_override["seed"] = _SEED.value
rl_override["runner.seed"] = _SEED.value
sim_backend = _SIM_BACKEND.value
rllib = None
policy_path = None
# Determine policy path: use explicit policy if provided, otherwise auto-discover
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:
policy_path = find_best_policy(env_name)
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 a policy using --policy flag or train a model first")
logger.error("Please specify --rllib or train a model first")
return
backend = get_inference_backend(policy_path)
backend = get_inference_backend(policy_path, rllib)
if backend == "jax":
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)
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

63
scripts/train.py
View File

@@ -31,12 +31,45 @@ _SIM_BACKEND = flags.DEFINE_string(
)
_NUM_ENVS = flags.DEFINE_integer("num-envs", 2048, "Number of envs to train")
_RENDER = flags.DEFINE_bool("render", False, "Render the env")
_TRAIN_BACKEND = flags.DEFINE_string("train-backend", "jax", "The learning backend. (jax/torch)")
_TRAIN_BACKEND = flags.DEFINE_string("train-backend", None, "The learning backend. (jax/torch)")
_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", "skrl", "The RL framework (skrl/rslrl)")
def get_train_backend(supports: utils.DeviceSupports):
def get_train_backend(supports: utils.DeviceSupports, train_backend_arg: str | None, rllib: str):
"""
Determine the training backend based on device supports, user input, and RL framework.
Args:
supports: Device support information
train_backend_arg: User-specified backend via --train-backend flag (None if not provided)
rllib: RL framework to use ("skrl" or "rslrl")
Returns:
The determined backend name ("jax" or "torch")
Raises:
Exception: If user specifies incompatible backend or no backend is available
"""
# RSLRL only supports PyTorch
if rllib == "rslrl":
if train_backend_arg is not None and train_backend_arg != "torch":
raise Exception("RSLRL only supports PyTorch backend.")
if not supports.torch:
raise Exception("RSLRL requires PyTorch, but it is not available on your device.")
return "torch"
# User explicitly specified backend
if train_backend_arg is not None:
backend = train_backend_arg
if backend == "jax" and not supports.jax:
raise Exception("JAX is not available on your device.")
if backend == "torch" and not supports.torch:
raise Exception("PyTorch is not available on your device.")
return backend
# Auto-select backend based on device priority
if supports.jax and supports.jax_gpu:
return "jax"
elif supports.torch and supports.torch_gpu:
@@ -46,7 +79,7 @@ def get_train_backend(supports: utils.DeviceSupports):
elif supports.torch:
return "torch"
else:
raise Exception("neither jax nor torch not avaliable on the device.")
raise Exception("Neither JAX nor PyTorch is available on the device.")
def main(argv):
@@ -61,19 +94,26 @@ def main(argv):
rl_override["num_envs"] = _NUM_ENVS.value
if _RAND_SEED.value:
rl_override["seed"] = None
rl_override["runner.seed"] = None
elif _SEED.present:
rl_override["seed"] = _SEED.value
rl_override["runner.seed"] = _SEED.value
sim_backend = _SIM_BACKEND.value
train_backend = "jax"
if not _TRAIN_BACKEND.present:
train_backend = get_train_backend(device_supports)
else:
train_backend = _TRAIN_BACKEND.value
rllib = _RLLIB.value
# Determine the training backend
train_backend = get_train_backend(device_supports, _TRAIN_BACKEND.value, rllib)
trainer = None
if train_backend == "jax":
if rllib == "rslrl":
# RSLRL training flow
assert device_supports.torch, "PyTorch is not available on your device"
assert train_backend == "torch", "RSLRL only supports PyTorch backend"
from motrix_rl.rslrl.torch.train import ppo
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
elif train_backend == "jax":
from motrix_rl.skrl.jax.train import ppo
config.jax.backend = "jax" # or "numpy"
@@ -82,7 +122,6 @@ def main(argv):
elif train_backend == "torch":
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)
else:
raise Exception(f"Unknown train backend: {train_backend}")