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

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}")