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

10
motrix_rl/pyproject.toml
View File

@@ -4,8 +4,8 @@ build-backend = "uv_build"
[project]
name = "motrix-rl"
version = "0.2.0"
description = "Reinforcement learning training framework based on SKRL with multi-backend unified training interface."
version = "0.3.0"
description = "Reinforcement learning training framework with multi-backend unified training interface."
authors = [{ name = "Motphys", email = "developers@motphys.com" }]
requires-python = "==3.10.*"
readme = "README.md"
@@ -25,6 +25,12 @@ skrl-torch = [
"torchvision===v0.22.0+cu128",
"torchaudio===2.7.0+cu128",
]
rslrl = [
"rsl-rl-lib>=4.0.0",
"torch===2.7.0+cu128",
"torchvision===v0.22.0+cu128",
"torchaudio===2.7.0+cu128",
]
test = ["pytest~=8.3.3"]
[tool.uv.sources]

8
motrix_rl/src/motrix_rl/__init__.py
View File

@@ -13,4 +13,10 @@
# limitations under the License.
# ==============================================================================
from . import cfgs, skrl # noqa: F401
from . import tasks # noqa: F401
from .rslrl.cfg import ( # noqa: F401
RslRlActorCfg,
RslRlCriticCfg,
RslRlPpoAlgorithmCfg,
RslrlRunnerCfg,
)

42
motrix_rl/src/motrix_rl/base.py
View File

@@ -1,42 +0,0 @@
# 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 dataclasses
from typing import Optional
@dataclasses.dataclass
class BaseRLCfg:
"""
Config for the reinforcement learning algorithm
"""
# Basic training parameters
seed: Optional[int] = None
num_envs: int = 2048
play_num_envs: int = 16
max_env_steps: int = 20480000
check_point_interval: int = 1000
def replace(self, **updates) -> "BaseRLCfg":
return dataclasses.replace(self, **updates)
@property
def max_batch_env_steps(self) -> int:
"""
The max batched environment steps for the RL algorithm.
"""
n = int(self.max_env_steps / self.num_envs)
return (int)(n / self.check_point_interval) * self.check_point_interval

553
motrix_rl/src/motrix_rl/cfgs.py
View File

@@ -1,553 +0,0 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.skrl.cfg import PPOCfg
class basic:
@rlcfg("cartpole")
@dataclass
class CartPolePPO(PPOCfg):
max_env_steps: int = 10_000_000
check_point_interval: int = 500
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32)
rollouts: int = 32
learning_epochs: int = 5
mini_batches: int = 4
@rlcfg("acrobot", backend="jax")
@dataclass
class AcrobotPPO(PPOCfg):
max_env_steps: int = 60_000_000
check_point_interval: int = 500
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32)
rollouts: int = 64
learning_epochs: int = 5
mini_batches: int = 8
learning_rate: float = 3e-4
grad_norm_clip: float = 0.5
clip_predicted_values: bool = False
value_clip: float = 10.0
entropy_loss_scale: float = 0.05
learning_rate_scheduler_kl_threshold: float = 0.02
discount_factor: float = 0.995
lambda_param: float = 0.97
ratio_clip: float = 0.2
value_loss_scale: float = 0.5
random_timesteps: int = 0
learning_starts: int = 0
kl_threshold: float = 0.03
grad_norm_clip: float = 0.1
entropy_loss_scale: float = 0.1
@rlcfg("acrobot", backend="torch")
@dataclass
class AcrobotPPOTorch(PPOCfg):
max_env_steps: int = 60_000_000
check_point_interval: int = 500
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32)
rollouts: int = 64
learning_epochs: int = 5
mini_batches: int = 8
learning_rate: float = 0.0003
learning_rate_scheduler_kl_threshold: float = 0.02
entropy_loss_scale: float = 0.2
discount_factor: float = 0.995
lambda_param: float = 0.97
@rlcfg("pendulum")
@dataclass
class PendulumPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 5_000_000
num_envs: int = 1024
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (64, 64)
value_hidden_layer_sizes: tuple[int, ...] = (64, 64)
learning_rate: float = 3e-4
rollouts: int = 32
learning_epochs: int = 5
mini_batches: int = 4
@rlcfg("bounce_ball")
@dataclass
class BounceBallPPO(PPOCfg):
max_env_steps: int = 100_000_000
check_point_interval: int = 5000
# Override PPO configuration for bounce ball task
policy_hidden_layer_sizes: tuple[int, ...] = (64, 64, 64)
value_hidden_layer_sizes: tuple[int, ...] = (64, 64, 64)
rollouts: int = 128
learning_epochs: int = 15
mini_batches: int = 16
learning_rate: float = 2e-4
num_envs: int = 1024
@rlcfg("dm-walker", backend="jax")
@rlcfg("dm-stander", backend="jax")
@rlcfg("dm-runner", backend="jax")
@dataclass
class WalkerPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 4
@rlcfg("dm-stander", backend="torch")
@rlcfg("dm-walker", backend="torch")
@dataclass
class WalkerPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
@rlcfg("dm-runner", backend="torch")
@dataclass
class RunnerPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 2
mini_batches: int = 32
@rlcfg("dm-cheetah", backend="jax")
@dataclass
class CheetahPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
@rlcfg("dm-cheetah", backend="torch")
@dataclass
class CheetahPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
@rlcfg("dm-hopper-stand", backend="jax")
@dataclass
class HopperStandPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 4
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
@rlcfg("dm-hopper-hop", backend="jax")
@dataclass
class HopperPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 5
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
@rlcfg("dm-hopper-stand", backend="torch")
@rlcfg("dm-hopper-hop", backend="torch")
@dataclass
class HopperPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 5
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
@rlcfg("dm-reacher", backend="jax")
@dataclass
class ReacherPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
value_hidden_layer_sizes: tuple[int, ...] = (32, 32, 32)
@rlcfg("dm-reacher", backend="torch")
@dataclass
class ReacherPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Override PPO configuration
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
policy_hidden_layer_sizes: tuple[int, ...] = (64, 64, 64)
value_hidden_layer_sizes: tuple[int, ...] = (64, 64, 64)
@rlcfg("dm-finger-spin")
@rlcfg("dm-finger-turn-easy")
@rlcfg("dm-finger-turn-hard")
@dataclass
class FingerPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Similar to other dm_* configs
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 4
# JAX-only overrides (keep torch + turn-easy unaffected)
@rlcfg("dm-finger-spin", backend="jax")
@dataclass
class FingerSpinPPOJax(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# More conservative PPO for stability (spin can collapse mid-training in JAX)
learning_rate: float = 7.5e-5
learning_rate_scheduler_kl_threshold: float = 0.003
entropy_loss_scale: float = 5e-4
rollouts: int = 24
learning_epochs: int = 1
mini_batches: int = 16
ratio_clip: float = 0.08
value_clip: float = 0.1
value_loss_scale: float = 0.5
grad_norm_clip: float = 0.25
@rlcfg("dm-finger-turn-hard", backend="jax")
@dataclass
class FingerTurnHardPPOJax(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
# Extra conservative to avoid late-stage collapses
learning_rate: float = 5e-5
learning_rate_scheduler_kl_threshold: float = 0.004
rollouts: int = 24
learning_epochs: int = 1
mini_batches: int = 16
ratio_clip: float = 0.08
value_loss_scale: float = 0.5
grad_norm_clip: float = 0.25
@rlcfg("dm-manipulator-bring-ball", backend="jax")
@dataclass
class ManipulatorPPOJax(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 20000
num_envs: int = 256
learning_rate: float = 3e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 4
ratio_clip: float = 0.2
entropy_loss_scale: float = 1e-3
grad_norm_clip: float = 1.0
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
@rlcfg("dm-manipulator-bring-ball", backend="torch")
@dataclass
class ManipulatorPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 20000
num_envs: int = 256
learning_rate: float = 2e-4
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 4
policy_hidden_layer_sizes: tuple[int, ...] = (256, 256)
value_hidden_layer_sizes: tuple[int, ...] = (256, 256)
@rlcfg("dm-humanoid-stand", backend="jax")
@rlcfg("dm-humanoid-walk", backend="jax")
@rlcfg("dm-humanoid-run", backend="jax")
@dataclass
class HumanoidPPO(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
learning_rate: float = 3e-4
rollouts: int = 24
learning_epochs: int = 8
mini_batches: int = 2
policy_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
value_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
@rlcfg("dm-humanoid-stand", backend="torch")
@rlcfg("dm-humanoid-walk", backend="torch")
@rlcfg("dm-humanoid-run", backend="torch")
@dataclass
class HumanoidPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 1024 * 40000
num_envs: int = 2048
learning_rate: float = 3e-4
rollouts: int = 24
learning_epochs: int = 8
mini_batches: int = 2
policy_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
value_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
class locomotion:
@rlcfg("go1-flat-terrain-walk")
@dataclass
class Go1WalkPPO(PPOCfg):
"""
Go1 Walk RL config
"""
seed: int = 42
share_policy_value_features: bool = False
max_env_steps: int = 1024 * 60_000
num_envs: int = 2048
# Override PPO configuration
rollouts: int = 24
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
learning_epochs: int = 5
mini_batches: int = 3
learning_rate: float = 3e-4
@rlcfg("go2-flat-terrain-walk")
@dataclass
class Go2WalkPPO(PPOCfg):
"""
Go2 Walk RL config
"""
seed: int = 42
share_policy_value_features: bool = False
max_env_steps: int = 1024 * 60_000
num_envs: int = 2048
# Override PPO configuration
rollouts: int = 24
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
learning_epochs: int = 5
mini_batches: int = 3
learning_rate: float = 3e-4
@rlcfg("go1-rough-terrain-walk")
@dataclass
class Go1WalkRoughPPO(Go1WalkPPO):
policy_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
value_hidden_layer_sizes: tuple[int, ...] = (512, 256, 128)
@rlcfg("go1-stairs-terrain-walk")
@dataclass
class Go1WalkStairsPPO(Go1WalkRoughPPO): ...
class manipulation:
@rlcfg("franka-lift-cube", "jax")
@dataclass
class FrankaLiftPPOJax(PPOCfg):
seed: int = 42
max_env_steps: int = 4096 * 50000
share_policy_value_features: bool = True
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
rollouts: int = 24
learning_epochs: int = 4
mini_batches: int = 32
learning_rate: float = 1e-3
@rlcfg("franka-lift-cube", "torch")
@dataclass
class FrankaLiftPPOTorch(PPOCfg):
seed: int = 42
max_env_steps: int = 4096 * 50000
share_policy_value_features: bool = True
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
rollouts: int = 24
learning_epochs: int = 8
mini_batches: int = 4
learning_rate: float = 3e-4
learning_rate_scheduler_kl_threshold: float = 0.01
entropy_loss_scale: float = 0.001
rewards_shaper_scale: float = 0.01
@rlcfg("franka-open-cabinet")
@dataclass
class FrankaOpenCabinetPPO(PPOCfg):
seed: int = 64
max_env_steps: int = 2048 * 24000
share_policy_value_features: bool = False
# Override PPO configuration
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
rollouts: int = 16
learning_epochs: int = 5
mini_batches: int = 32
learning_rate: float = 3e-4
rewards_shaper_scale: float = 1e-1
entropy_loss_scale: float = 0.001
@rlcfg("shadow-hand-repose")
@dataclass
class ShadowHandReposePPO(PPOCfg):
"""
Shadow Hand Repose PPO configuration
"""
# ===== Basic Settings =====
seed: int = 42
num_envs: int = 8192
play_num_envs: int = 16
max_env_steps: int = 200_000_000
check_point_interval: int = 1000
# ===== Network Architecture =====
policy_hidden_layer_sizes: tuple[int, ...] = (512, 512, 256, 128)
value_hidden_layer_sizes: tuple[int, ...] = (512, 512, 256, 128)
share_policy_value_features: bool = True
# ===== PPO Core Parameters =====
rollouts: int = 16
learning_epochs: int = 5
mini_batches: int = 4
discount_factor: 0.99
lambda_param: 0.95
# ===== Learning Rate =====
learning_rate: float = 5.0e-04
learning_rate_scheduler_kl_threshold: float = 0.016
# ===== Clipping =====
ratio_clip: float = 0.2
value_clip: float = 0.2
clip_predicted_values: bool = True
grad_norm_clip: float = 1.0
# ===== Loss Coefficients =====
entropy_loss_scale: float = 0.0
value_loss_scale: float = 2.0
kl_threshold: float = 0.0
# ===== Reward Shaping =====
rewards_shaper_scale: float = 0.01
# ===== Training Control =====
random_timesteps: int = 0
learning_starts: int = 0
time_limit_bootstrap: bool = False
class navigation:
@rlcfg("anymal_c_navigation_flat")
@dataclass
class AnymalCPPOConfig(PPOCfg):
# ===== Basic Training Parameters =====
seed: int = 42 # Random seed
num_envs: int = 2048 # Number of parallel environments during training
play_num_envs: int = 16 # Number of parallel environments during evaluation
max_env_steps: int = 100_000_000 # Maximum training steps
# ===== PPO Algorithm Core Parameters =====
learning_rate: float = 3e-4 # Learning rate
rollouts: int = 48 # Number of experience replay rollouts
learning_epochs: int = 6 # Number of training epochs per update
mini_batches: int = 32 # Number of mini-batches
discount_factor: float = 0.99 # Discount factor
lambda_param: float = 0.95 # GAE parameter
grad_norm_clip: float = 1.0 # Gradient clipping
# ===== PPO Clipping Parameters =====
ratio_clip: float = 0.2 # PPO clipping ratio
value_clip: float = 0.2 # Value clipping
clip_predicted_values: bool = True # Clip predicted values
# Medium-sized network (default configuration, suitable for most tasks)
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)

48
motrix_rl/src/motrix_rl/registry.py
View File

@@ -15,19 +15,18 @@
import logging
from dataclasses import dataclass, field
from typing import Callable, Type, TypeVar
from typing import Any, Callable, Type, TypeVar
from motrix_envs import registry as env_registry
from motrix_rl.base import BaseRLCfg
logger = logging.getLogger(__name__)
TRLCfg = TypeVar("TRLCfg", bound=BaseRLCfg)
TRLCfg = TypeVar("TRLCfg")
@dataclass
class EnvRlCfgs:
cfgs: dict[str, dict[str, Type[BaseRLCfg]]] = field(default_factory=dict)
cfgs: dict[str, dict[str, Type]] = field(default_factory=dict)
"""
The RL configuration classes available for this environment.
Structure: {rl_framework: {backend: config_class}}
@@ -39,7 +38,7 @@ class EnvRlCfgs:
_rlcfgs: dict[str, EnvRlCfgs] = {}
def _register_rlcfg(env_name: str, rllib: str, backend: str, train_cfg_cls: Type[BaseRLCfg]):
def _register_rlcfg(env_name: str, rllib: str, backend: str, train_cfg_cls: Type):
"""
Register a training configuration class for an environment, reinforcement learning framework, and backend.
@@ -60,9 +59,39 @@ def _register_rlcfg(env_name: str, rllib: str, backend: str, train_cfg_cls: Type
_rlcfgs[env_name].cfgs[rllib][backend] = train_cfg_cls
def _infer_framework_from_class(cls: Type) -> str:
"""Infer RL framework name from the class's parent class.
Args:
cls: Configuration class to inspect
Returns:
Framework name ("skrl" or "rslrl")
Raises:
ValueError: If framework cannot be determined from parent class
"""
# Import here to avoid circular imports
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
# Check entire MRO (Method Resolution Order) for framework base classes
for base in cls.__mro__:
if base is SkrlCfg:
return "skrl"
elif base is RslrlCfg:
return "rslrl"
raise ValueError(
f"Cannot infer RL framework from {cls.__name__}. Class must inherit from either SkrlCfg or RslrlCfg."
)
def rlcfg(env_name: str, backend: str = None) -> Callable[[Type[TRLCfg]], Type[TRLCfg]]:
"""
Decorator to register a training configuration class for an environment, RL framework, and backend.
Decorator to register a training configuration class for an environment and backend.
The RL framework (skrl/rslrl) is automatically inferred from the parent class.
Args:
env_name: Environment name
@@ -70,15 +99,18 @@ def rlcfg(env_name: str, backend: str = None) -> Callable[[Type[TRLCfg]], Type[T
"""
def decorator(cls: Type[TRLCfg]) -> Type[TRLCfg]:
# Infer framework from parent class
rl_framework = _infer_framework_from_class(cls)
backends = ["jax", "torch"] if backend is None else [backend]
for b in backends:
_register_rlcfg(env_name, "skrl", b, cls)
_register_rlcfg(env_name, rl_framework, b, cls)
return cls
return decorator
def default_rl_cfg(env_name: str, rllib: str, backend: str) -> BaseRLCfg:
def default_rl_cfg(env_name: str, rllib: str, backend: str) -> Any:
"""
Get the default training configuration for an environment, reinforcement learning framework, and backend.

40
motrix_rl/src/motrix_rl/rslrl/__init__.py Normal file
View File

@@ -0,0 +1,40 @@
# 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.
# ==============================================================================
"""RSLRL integration module for MotrixLab.
This module provides configuration classes and utilities for using RSLRL
(ETH Zurich's RL library) with MotrixLab.
The configuration structure matches rsl_rl's flat format with separate
actor and critic configs at the top level.
"""
from motrix_rl.rslrl.cfg import (
RslRlActorCfg,
RslRlCriticCfg,
RslRlPpoAlgorithmCfg,
RslrlRunnerCfg,
)
__all__ = [
"RslRlActorCfg",
"RslRlCriticCfg",
"RslRlPpoAlgorithmCfg",
"RslrlRunnerCfg",
"field_override",
"inherit_field",
"configclass",
]

122
motrix_rl/src/motrix_rl/rslrl/cfg.py Normal file
View File

@@ -0,0 +1,122 @@
# 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.
# ==============================================================================
"""PPO Agent Configuration for RSLRL
This module provides configuration classes for PPO agents using the
RSLRL (ETH Zurich RL library) framework.
The configuration structure matches rsl_rl's flat format with separate
actor and critic configs at the top level.
"""
from dataclasses import dataclass, field
from typing import Literal
from motrix_rl.utils import class_to_dict
@dataclass
class RslRlActorCfg:
"""Configuration for the actor network."""
class_name: str = "MLPModel"
hidden_dims: list[int] = field(default_factory=lambda: [256, 128, 64])
activation: str = "elu"
obs_normalization: bool = True
stochastic: bool = True
init_noise_std: float = 1.0
noise_std_type: Literal["scalar", "log"] = "scalar"
state_dependent_std: bool = False
@dataclass
class RslRlCriticCfg:
"""Configuration for the critic network."""
class_name: str = "MLPModel"
hidden_dims: list[int] = field(default_factory=lambda: [256, 128, 64])
activation: str = "elu"
obs_normalization: bool = True
stochastic: bool = False
@dataclass
class RslRlPpoAlgorithmCfg:
"""Configuration for the PPO algorithm."""
class_name: str = "PPO"
optimizer: str = "adam"
learning_rate: float = 3e-4
num_learning_epochs: int = 5
num_mini_batches: int = 4
schedule: str = "adaptive"
value_loss_coef: float = 1.0
clip_param: float = 0.2
use_clipped_value_loss: bool = True
desired_kl: float = 0.008
entropy_coef: float = 0.01
gamma: float = 0.99
lam: float = 0.95
max_grad_norm: float = 1.0
normalize_advantage_per_mini_batch: bool = False
rnd_cfg: dict | None = None
symmetry_cfg: dict | None = None
@dataclass
class RslrlRunnerCfg:
"""Configuration matching rsl_rl's flat structure.
This configuration provides separate actor and critic configs at the top level,
matching the structure expected by rsl_rl's OnPolicyRunner.
"""
# Runner settings
class_name: str = "OnPolicyRunner"
seed: int = 42
device: str = "cuda:0"
num_steps_per_env: int = 24
max_iterations: int = 10000
save_interval: int = 50
experiment_name: str = "experiment"
run_name: str = ""
logger: Literal["tensorboard", "neptune", "wandb"] = "tensorboard"
# Observation groups
obs_groups: dict[str, list[str]] = field(default_factory=lambda: {"actor": ["policy"], "critic": ["policy"]})
# Network configs - TOP LEVEL
actor: RslRlActorCfg = field(default_factory=RslRlActorCfg)
critic: RslRlCriticCfg = field(default_factory=RslRlCriticCfg)
algorithm: RslRlPpoAlgorithmCfg = field(default_factory=RslRlPpoAlgorithmCfg)
def to_dict(self) -> dict:
"""Convert config to dictionary for OnPolicyRunner.
Returns:
Dictionary representation matching rsl_rl's expected format.
"""
return class_to_dict(self)
@dataclass
class RslrlCfg:
"""Configuration for RSLRL."""
num_envs: int = 2048
play_num_envs: int = 16
runner: RslrlRunnerCfg = field(default_factory=RslrlRunnerCfg)

20
motrix_rl/src/motrix_rl/rslrl/torch/__init__.py Normal file
View File

@@ -0,0 +1,20 @@
# 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.
# ==============================================================================
"""PyTorch backend for RSLRL integration."""
from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap
__all__ = ["RslrlNpEnvWrap"]

20
motrix_rl/src/motrix_rl/rslrl/torch/train/__init__.py Normal file
View File

@@ -0,0 +1,20 @@
# 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.
# ==============================================================================
"""Training module for RSLRL integration."""
from motrix_rl.rslrl.torch.train.ppo import Trainer
__all__ = ["Trainer"]

172
motrix_rl/src/motrix_rl/rslrl/torch/train/ppo.py Normal file
View File

@@ -0,0 +1,172 @@
# 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.
# ==============================================================================
"""PPO Trainer for RSLRL integration."""
import logging
import torch
from rsl_rl.runners import OnPolicyRunner
from motrix_envs import registry as env_registry
from motrix_rl import registry as rl_registry
from motrix_rl import utils
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap
from motrix_rl.skrl import get_log_dir
logger = logging.getLogger(__name__)
class Trainer:
"""RSLRL PPO Trainer.
This class wraps RSLRL's OnPolicyRunner to provide a training interface
consistent with the SKRL trainer implementation.
"""
_env_name: str
_sim_backend: str
_rlcfg: RslrlCfg
_enable_render: bool
def __init__(
self,
env_name: str,
sim_backend: str = None,
enable_render: bool = False,
cfg_override: dict = None,
) -> None:
"""Initialize the RSLRL PPO trainer.
Args:
env_name: Name of the environment to train
sim_backend: Simulation backend to use (e.g., "mujoco", "npcm")
enable_render: Whether to enable rendering during training
cfg_override: Optional configuration overrides
"""
rlcfg = rl_registry.default_rl_cfg(env_name, "rslrl", backend="torch")
if cfg_override is not None:
rlcfg = utils.cfg_override(rlcfg, cfg_override)
self._rlcfg = rlcfg
self._env_name = env_name
self._sim_backend = sim_backend
self._enable_render = enable_render
def train(self) -> None:
"""Start training the agent.
Creates the environment, wraps it for RSLRL, and runs the training loop.
"""
rlcfg = self._rlcfg
# Create environment
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
# Set random seed
if rlcfg.runner.seed is not None:
torch.manual_seed(rlcfg.runner.seed)
# Determine device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
# Wrap environment for RSLRL
vec_env = RslrlNpEnvWrap(env, device)
# Create RSLRL config - use to_dict() method
rslrl_cfg = self._create_rslrl_config()
# Create RSLRL runner
runner = OnPolicyRunner(
vec_env, rslrl_cfg, log_dir=get_log_dir(self._env_name, rllib="rslrl", agent_name="PPO"), device=device
)
# Start training
logger.info(f"Starting training for {self._env_name}")
logger.info(f"Number of environments: {rlcfg.num_envs}")
# Get max_iterations from config
total_iterations = rslrl_cfg["max_iterations"]
logger.info(f"Number of learning iterations: {total_iterations}")
runner.learn(num_learning_iterations=total_iterations)
logger.info("Training completed")
def play(self, policy_path: str) -> None:
"""Evaluate a trained policy.
Args:
policy_path: Path to the saved policy file
"""
import time
rlcfg = self._rlcfg
# Create environment with play_num_envs
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
# Set random seed
if rlcfg.runner.seed is not None:
torch.manual_seed(rlcfg.runner.seed)
# Determine device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Wrap environment for RSLRL
vec_env = RslrlNpEnvWrap(env, device)
# Create RSLRL config (minimal for evaluation)
rslrl_cfg = self._create_rslrl_config()
# Create RSLRL runner with log_dir=None to disable logging (no git diff storage in play mode)
runner = OnPolicyRunner(vec_env, rslrl_cfg, log_dir=None, device=device)
# Load policy
logger.info(f"Loading policy from {policy_path}")
runner.load(policy_path)
# Run evaluation loop
logger.info("Starting evaluation loop...")
logger.info("Press Ctrl+C to stop")
obs, _ = vec_env.reset()
fps = 60
try:
while True:
t = time.time()
# Get actions from policy
with torch.no_grad():
policy = runner.get_inference_policy(device=device)
# MLPModel is callable, returns distribution mean for deterministic evaluation
actions = policy(obs)
# Step environment
obs, rewards, dones, infos = vec_env.step(actions)
# Render the environment
vec_env.render()
delta_time = time.time() - t
if delta_time < 1.0 / fps:
time.sleep(1.0 / fps - delta_time)
except KeyboardInterrupt:
logger.info("Evaluation interrupted by user")
def _create_rslrl_config(self) -> dict:
return self._rlcfg.runner.to_dict()

168
motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py Normal file
View File

@@ -0,0 +1,168 @@
# 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.
# ==============================================================================
"""VecEnv wrapper for adapting NpEnv to RSLRL's VecEnv interface."""
import numpy as np
import torch
from rsl_rl.env.vec_env import VecEnv
from tensordict import TensorDict
from motrix_envs.np.env import NpEnv
class RslrlNpEnvWrap(VecEnv):
"""Adapter class that wraps NpEnv to RSLRL's VecEnv interface.
RSLRL expects a VecEnv interface with specific methods for stepping,
resetting, and accessing observations. This adapter converts between
NpEnv's NpEnvState format and RSLRL's expected format.
"""
def __init__(self, env: NpEnv, device: torch.device):
"""Initialize the VecEnv adapter.
Args:
env: The NpEnv instance to wrap
device: PyTorch device for tensors
"""
self._env = env
self._device = device
self._state = None
self._num_envs = env.num_envs
self._viewer = None # Will be initialized lazily when render() is called
# Set max_episode_length from env config
self.max_episode_length = self._env.cfg.max_episode_steps if self._env.cfg.max_episode_steps else 10000
# Episode length buffer for tracking
self.episode_length_buf = torch.zeros(self._num_envs, dtype=torch.long, device=self._device)
# Configuration dict for RSLRL logger
self.cfg = {
"env_name": self._env.cfg.__class__.__name__,
}
# Initialize the environment state
self.reset()
@property
def num_envs(self) -> int:
"""Number of parallel environments."""
return self._num_envs
@property
def num_obs(self) -> int:
"""Size of observation space."""
return self._env.observation_space.shape[0]
@property
def num_actions(self) -> int:
"""Size of action space."""
return self._env.action_space.shape[0]
@property
def device(self) -> torch.device:
"""PyTorch device for tensors."""
return self._device
@property
def unwrapped(self) -> "RslrlNpEnvWrap":
"""Return the unwrapped environment (self for this wrapper)."""
return self
def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
# Convert torch actions to numpy
actions_np = actions.cpu().numpy()
# Step the environment
state = self._env.step(actions_np)
self._state = state
# Update episode length buffer
self.episode_length_buf += 1
# Reset episode length for done environments
dones_np = state.done.astype(bool)
self.episode_length_buf[dones_np] = 0
# Convert to torch tensors
obs_tensor = torch.from_numpy(state.obs).to(self._device)
rewards = torch.from_numpy(state.reward).to(self._device)
# Merge terminated and truncated into dones
dones = torch.from_numpy(state.done.astype(np.float32)).to(self._device)
# Create TensorDict for observations
obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
# Build extras dict (RSLRL calls it "extras" not "infos")
extras = {}
if "time_outs" in state.info:
extras["time_outs"] = torch.from_numpy(state.info["time_outs"]).to(self._device)
return obs, rewards, dones, extras
def reset(self) -> tuple[TensorDict, dict]:
"""Reset all environments.
Returns:
Tuple of (observations, extras)
- observations: TensorDict with observation groups
- extras: dict with episode information
"""
state = self._env.init_state()
self._state = state
# Reset episode length buffer
self.episode_length_buf.zero_()
obs_tensor = torch.from_numpy(state.obs).to(self._device)
# Create TensorDict for observations
obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
# Build extras dict
extras = {}
return obs, extras
def get_observations(self) -> TensorDict:
"""Get current observations without stepping the environment.
Returns:
Current observations as TensorDict
"""
if self._state is None:
obs, _ = self.reset()
return obs
obs_tensor = torch.from_numpy(self._state.obs).to(self._device)
obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
return obs
def render(self) -> None:
"""Render the environment.
For NpEnv, this triggers the motrixsim viewer to display the scene.
The viewer window must be kept alive by calling this method regularly.
"""
# Initialize viewer on first call
if self._viewer is None:
from motrix_envs.np.renderer import NpRenderer
self._viewer = NpRenderer(env=self._env)
# Render the current state
self._viewer.render()

28
motrix_rl/src/motrix_rl/skrl/__init__.py
View File

@@ -13,9 +13,31 @@
# limitations under the License.
# ==============================================================================
from datetime import datetime
LOG_DIR_PREFIX = "runs"
def get_log_dir(env_name: str) -> str:
"""Get the log directory for the given environment name."""
return f"{LOG_DIR_PREFIX}/{env_name}"
def get_log_dir(env_name: str, rllib: str = "skrl", agent_name: str = "PPO") -> str:
"""Get the log directory for the given environment name and RL framework.
Args:
env_name: Name of the environment
rllib: RL framework name (e.g., "skrl", "rslrl")
agent_name: Name of the agent (e.g., "PPO")
Returns:
Log directory path:
- For SKRL: runs/{env_name}/{rllib}/ (SKRL creates its own timestamp subdirectory)
- For RSLRL: runs/{env_name}/{rllib}/{time}_{agent}/ (RSLRL doesn't create subdirectories)
"""
if rllib == "skrl":
# SKRL creates its own timestamp subdirectory, so we just provide the base path
return f"{LOG_DIR_PREFIX}/{env_name}/{rllib}"
else:
# RSLRL doesn't create subdirectories, so we add the timestamp here
now = datetime.now()
time_str = now.strftime("%y-%m-%d_%H-%M-%S")
microseconds = now.microsecond
time_str = f"{time_str}-_{microseconds:05d}"
return f"{LOG_DIR_PREFIX}/{env_name}/{rllib}/{time_str}_{agent_name}"

73
motrix_rl/src/motrix_rl/skrl/cfg.py
View File

@@ -1,73 +0,0 @@
# 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.
# ==============================================================================
"""PPO Agent Configuration for SKRL
This module provides configuration classes for PPO agents that match the
YAML configuration structure used in SKRL.
"""
from dataclasses import dataclass
from motrix_rl.base import BaseRLCfg
@dataclass
class PPOCfg(BaseRLCfg):
"""PPO configuration .
This class provides all the parameters needed to configure a PPO agent
in SKRL
"""
# Model architecture settings
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
# Whether to share feature extraction layers between policy and value networks. only works if:
# 1. both networks have the same architecture
# 2. the backend is torch
share_policy_value_features: bool = True
# Agent settings
rollouts: int = 32
learning_epochs: int = 2
mini_batches: int = 32
discount_factor: float = 0.99
lambda_param: float = 0.95
# Learning rate settings
learning_rate: float = 1e-3
learning_rate_scheduler_kl_threshold: float = 0.008
# Training settings
random_timesteps: int = 0
learning_starts: int = 0
grad_norm_clip: float = 1.0
time_limit_bootstrap: bool = True
# PPO clipping settings
ratio_clip: float = 0.2
value_clip: float = 0.2
clip_predicted_values: bool = True
# Loss settings
entropy_loss_scale: float = 0.0
value_loss_scale: float = 2.0
kl_threshold: float = 0
# Reward shaping
rewards_shaper_scale: float = 1.0

466
motrix_rl/src/motrix_rl/skrl/config.py Normal file
View File

@@ -0,0 +1,466 @@
# 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.
# ==============================================================================
"""SKRL Configuration Classes
This module provides configuration classes for SKRL RL framework integration.
The configuration structure matches template/skrl_config.yaml with a hierarchical
design following the same pattern as RslrlConfig.
Configuration Hierarchy:
SkrlCfg (top-level) -> SkrlRunnerCfg (runner-level)
-> SkrlModelsCfg (models)
-> SkrlPolicyCfg (policy model)
-> SkrlValueCfg (value model)
-> SkrlMemoryCfg (memory)
-> SkrlAgentCfg (PPO agent)
-> SkrlAgentExperimentCfg (experiment settings)
-> SkrlTrainerCfg (trainer)
"""
import dataclasses
from dataclasses import dataclass, field
@dataclass
class SkrlPolicyCfg:
"""Configuration for SKRL policy (GaussianMixin) model.
Corresponds to the policy section in template/skrl_config.yaml.
"""
class_name: str = "GaussianMixin"
clip_actions: bool = False
clip_log_std: bool = True
initial_log_std: float = 1.0
min_log_std: float = -20.0
max_log_std: float = 2.0
reduction: str = "sum"
input: str = "STATES"
hiddens: list[int] = field(default_factory=lambda: [256, 128, 64])
hidden_activation: list[str] = field(default_factory=lambda: ["elu"])
output: str = "ACTIONS"
output_activation: str = ""
output_scale: float = 1.0
def _normalize_activations(self, num_layers: int) -> str | list[str]:
"""Normalize hidden_activation to match num_layers.
SKRL requires either a single activation string (applied to all layers)
or a list with length matching the number of layers.
Args:
num_layers: Number of hidden layers (len(self.hiddens))
Returns:
str or list[str] suitable for SKRL's network format
Raises:
ValueError: If activation list length > 1 and doesn't match num_layers
"""
activations = self.hidden_activation
# Empty list -> no activations
if isinstance(activations, list) and len(activations) == 0:
return [""] * num_layers
# Single element list -> convert to string (SKRL will replicate)
if isinstance(activations, list) and len(activations) == 1:
return activations[0]
# String -> return as-is (SKRL will replicate)
if isinstance(activations, str):
return activations
# List with matching length -> use as-is
if isinstance(activations, list) and len(activations) == num_layers:
return activations
# List with mismatched length > 1 -> raise error
if isinstance(activations, list) and len(activations) > 1:
raise ValueError(
f"hidden_activation length ({len(activations)}) must match "
f"the number of hidden layers ({num_layers}), or be a single value "
f"to apply to all layers. Got hiddens={self.hiddens}, "
f"hidden_activation={activations}"
)
return activations
def to_network(self) -> tuple[list[dict], str]:
"""Convert configuration to SKRL's network and output format.
Returns:
(network, output) tuple where:
- network: SKRL network definition list of dicts
- output: SKRL output expression string (e.g., "tanh(ACTIONS)", "ONE")
Examples:
Policy with hiddens=[256,128,64], output_activation="tanh", output_scale=1.0:
network = [{"name": "net", "input": "STATES", "layers": [256,128,64], "activations": "elu"}]
output = "tanh(ACTIONS)"
Value with hiddens=[256,128,64], output_activation="", output_scale=0.5:
network = [{"name": "net", "input": "STATES", "layers": [256,128,64], "activations": "elu"}]
output = "0.5 * ONE"
"""
# Normalize activations to match hiddens length
num_layers = len(self.hiddens)
activations = self._normalize_activations(num_layers)
# Build network definition
network = [
{
"name": "net",
"input": "STATES",
"layers": self.hiddens,
"activations": activations,
}
]
# Build output expression
# Use output field directly (already in correct format)
# Apply scale if not 1.0
scale_prefix = f"{self.output_scale} * " if self.output_scale != 1.0 else ""
# Apply activation if specified
if self.output_activation:
output = f"{scale_prefix}{self.output_activation}({self.output})"
else:
output = f"{scale_prefix}{self.output}"
return network, output
def to_dict(self) -> dict:
"""Convert to dict, mapping class_name to class."""
from motrix_rl.utils import class_to_dict
return class_to_dict(self)
@dataclass
class SkrlValueCfg:
"""Configuration for SKRL value (DeterministicMixin) model.
Corresponds to the value section in template/skrl_config.yaml.
"""
class_name: str = "DeterministicMixin"
clip_actions: bool = False
input: str = "STATES"
hiddens: list[int] = field(default_factory=lambda: [256, 128, 64])
hidden_activation: list[str] = field(default_factory=lambda: ["elu"])
output: str = "ONE"
output_activation: str = ""
output_scale: float = 1.0
def _normalize_activations(self, num_layers: int) -> str | list[str]:
"""Normalize hidden_activation to match num_layers.
SKRL requires either a single activation string (applied to all layers)
or a list with length matching the number of layers.
Args:
num_layers: Number of hidden layers (len(self.hiddens))
Returns:
str or list[str] suitable for SKRL's network format
Raises:
ValueError: If activation list length > 1 and doesn't match num_layers
"""
activations = self.hidden_activation
# Empty list -> no activations
if isinstance(activations, list) and len(activations) == 0:
return [""] * num_layers
# Single element list -> convert to string (SKRL will replicate)
if isinstance(activations, list) and len(activations) == 1:
return activations[0]
# String -> return as-is (SKRL will replicate)
if isinstance(activations, str):
return activations
# List with matching length -> use as-is
if isinstance(activations, list) and len(activations) == num_layers:
return activations
# List with mismatched length > 1 -> raise error
if isinstance(activations, list) and len(activations) > 1:
raise ValueError(
f"hidden_activation length ({len(activations)}) must match "
f"the number of hidden layers ({num_layers}), or be a single value "
f"to apply to all layers. Got hiddens={self.hiddens}, "
f"hidden_activation={activations}"
)
return activations
def to_network(self) -> tuple[list[dict], str]:
"""Convert configuration to SKRL's network and output format.
Returns:
(network, output) tuple where:
- network: SKRL network definition list of dicts
- output: SKRL output expression string (e.g., "tanh(ACTIONS)", "ONE")
Examples:
Policy with hiddens=[256,128,64], output_activation="tanh", output_scale=1.0:
network = [{"name": "net", "input": "STATES", "layers": [256,128,64], "activations": "elu"}]
output = "tanh(ACTIONS)"
Value with hiddens=[256,128,64], output_activation="", output_scale=0.5:
network = [{"name": "net", "input": "STATES", "layers": [256,128,64], "activations": "elu"}]
output = "0.5 * ONE"
"""
# Normalize activations to match hiddens length
num_layers = len(self.hiddens)
activations = self._normalize_activations(num_layers)
# Build network definition
network = [
{
"name": "net",
"input": "STATES",
"layers": self.hiddens,
"activations": activations,
}
]
# Build output expression
# Use output field directly (already in correct format)
# Apply scale if not 1.0
scale_prefix = f"{self.output_scale} * " if self.output_scale != 1.0 else ""
# Apply activation if specified
if self.output_activation:
output = f"{scale_prefix}{self.output_activation}({self.output})"
else:
output = f"{scale_prefix}{self.output}"
return network, output
def to_dict(self) -> dict:
"""Convert to dict, mapping class_name to class."""
from motrix_rl.utils import class_to_dict
return class_to_dict(self)
@dataclass
class SkrlModelsCfg:
"""Configuration for SKRL models section.
Corresponds to the models section in template/skrl_config.yaml.
"""
separate: bool = False
policy: SkrlPolicyCfg = field(default_factory=SkrlPolicyCfg)
value: SkrlValueCfg = field(default_factory=SkrlValueCfg)
def to_dict(self) -> dict:
"""Convert to dict with nested configs."""
return {
"separate": self.separate,
"policy": self.policy.to_dict(),
"value": self.value.to_dict(),
}
@dataclass
class SkrlMemoryCfg:
"""Configuration for SKRL memory.
Corresponds to the memory section in template/skrl_config.yaml.
"""
class_name: str = "RandomMemory"
memory_size: int = -1 # -1: automatically determined
def to_dict(self) -> dict:
"""Convert to dict, mapping class_name to class."""
from motrix_rl.utils import class_to_dict
return class_to_dict(self)
@dataclass
class SkrlAgentExperimentCfg:
"""Experiment settings within agent config.
Corresponds to the experiment subsection in template/skrl_config.yaml.
"""
directory: str = "runs"
experiment_name: str = ""
write_interval: int = -1
checkpoint_interval: int = -1
@dataclass
class SkrlAgentCfg:
"""Configuration for SKRL PPO agent.
Corresponds to the agent section in template/skrl_config.yaml.
Field names match PPO_DEFAULT_CONFIG from SKRL.
"""
class_name: str = "PPO"
rollouts: int = 32
learning_epochs: int = 2
mini_batches: int = 32
discount_factor: float = 0.99
lam: float = 0.95
learning_rate: float = 1e-3
learning_rate_scheduler: str = "KLAdaptiveLR"
learning_rate_scheduler_kwargs: dict = field(default_factory=lambda: {"kl_threshold": 0.008})
random_timesteps: int = 0
learning_starts: int = 0
grad_norm_clip: float = 1.0
ratio_clip: float = 0.2
value_clip: float = 0.2
clip_predicted_values: bool = True
entropy_loss_scale: float = 0.0
value_loss_scale: float = 2.0
kl_threshold: int = 0
rewards_shaper_scale: float = 1.0
time_limit_bootstrap: bool = True
experiment: SkrlAgentExperimentCfg = field(default_factory=SkrlAgentExperimentCfg)
def to_dict(self) -> dict:
"""Convert configuration to dictionary for SKRL PPO agent.
Returns:
Dictionary representation matching SKRL's PPO agent configuration format.
Maps class_name -> class and lam -> lambda for SKRL compatibility.
Note:
- Maps 'class_name' to 'class' (SKRL convention)
- Maps 'lam' to 'lambda' (Python keyword conflict)
- Converts nested experiment config to dict
- Excludes state/value preprocessor fields (added dynamically during training)
"""
# Build base configuration dict
result = {
"class": self.class_name,
"rollouts": self.rollouts,
"learning_epochs": self.learning_epochs,
"mini_batches": self.mini_batches,
"discount_factor": self.discount_factor,
"lambda": self.lam,
"learning_rate": self.learning_rate,
"learning_rate_scheduler": self.learning_rate_scheduler,
"learning_rate_scheduler_kwargs": self.learning_rate_scheduler_kwargs,
"random_timesteps": self.random_timesteps,
"learning_starts": self.learning_starts,
"grad_norm_clip": self.grad_norm_clip,
"ratio_clip": self.ratio_clip,
"value_clip": self.value_clip,
"clip_predicted_values": self.clip_predicted_values,
"entropy_loss_scale": self.entropy_loss_scale,
"value_loss_scale": self.value_loss_scale,
"kl_threshold": self.kl_threshold,
"rewards_shaper_scale": self.rewards_shaper_scale,
"time_limit_bootstrap": self.time_limit_bootstrap,
"experiment": {
"directory": self.experiment.directory,
"experiment_name": self.experiment.experiment_name,
"write_interval": self.experiment.write_interval,
"checkpoint_interval": self.experiment.checkpoint_interval,
},
}
return result
@dataclass
class SkrlTrainerCfg:
"""Configuration for SKRL sequential trainer.
Corresponds to the trainer section in template/skrl_config.yaml.
"""
class_name: str = "SequentialTrainer"
timesteps: int = 10000
"""
The max number of batch env steps to run
"""
def to_dict(self) -> dict:
"""Convert to dict, mapping class_name to class."""
from motrix_rl.utils import class_to_dict
return class_to_dict(self)
@dataclass
class SkrlRunnerCfg:
"""Main SKRL runner configuration.
This mirrors the structure in template/skrl_config.yaml.
Follows the same pattern as RslrlRunnerCfg with nested configs
and a to_dict() method for dictionary conversion.
"""
seed: int = 42
models: SkrlModelsCfg = field(default_factory=SkrlModelsCfg)
memory: SkrlMemoryCfg = field(default_factory=SkrlMemoryCfg)
agent: SkrlAgentCfg = field(default_factory=SkrlAgentCfg)
trainer: SkrlTrainerCfg = field(default_factory=SkrlTrainerCfg)
def to_dict(self) -> dict:
"""Convert config to dictionary for SKRL.
Returns:
Dictionary representation matching SKRL's expected format.
Maps class_name -> class for all nested configs.
Note:
This method ensures that the output dictionary matches the exact
structure of template/skrl_config.yaml, including the 'class' field
names (instead of 'class_name' used in Python to avoid keyword conflicts).
"""
result = {
"seed": self.seed,
"models": self.models.to_dict(),
"memory": self.memory.to_dict(),
"agent": self.agent.to_dict(),
"trainer": self.trainer.to_dict(),
}
return result
@dataclass
class SkrlCfg:
"""Top-level SKRL configuration.
Follows the same pattern as RslrlCfg with environment-level settings
at the top level and runner configuration nested.
"""
# Basic training parameters
num_envs: int = 2048
play_num_envs: int = 16
runner: SkrlRunnerCfg = field(default_factory=SkrlRunnerCfg)
def replace(self, **updates) -> "SkrlCfg":
"""Replace specified fields and return a new instance."""
return dataclasses.replace(self, **updates)

209
motrix_rl/src/motrix_rl/skrl/jax/train/ppo.py
View File

@@ -19,9 +19,7 @@ import flax.linen as nn
import jax.numpy as jnp
import numpy as np
from skrl.agents.jax.ppo import PPO as BasePPO
from skrl.agents.jax.ppo import PPO_DEFAULT_CONFIG
from skrl.envs.jax import Wrapper
from skrl.memories.jax import RandomMemory
from skrl.models.jax import DeterministicMixin, GaussianMixin, Model
from skrl.resources.preprocessors.jax import RunningStandardScaler
from skrl.resources.schedulers.jax import KLAdaptiveRL
@@ -29,42 +27,69 @@ from skrl.trainers.jax import SequentialTrainer
from skrl.utils import set_seed
from motrix_envs import registry as env_registry
from motrix_rl import registry
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir
from motrix_rl.skrl.cfg import PPOCfg
from motrix_rl.skrl.config import SkrlCfg, SkrlMemoryCfg
from motrix_rl.skrl.jax import wrap_env
def _get_cfg(
rlcfg: PPOCfg,
def _instantiate_memory(memory_cfg: SkrlMemoryCfg, memory_size: int, num_envs: int, device) -> Any:
"""Instantiate a SKRL Memory class based on configuration.
Args:
memory_cfg: Memory configuration with class_name and settings
memory_size: Size of the memory buffer
num_envs: Number of parallel environments
device: Device to place memory on
Returns:
Instantiated SKRL Memory object
Raises:
ValueError: If class_name is not supported
"""
from skrl.memories.jax import RandomMemory
# Map class_name to actual Memory class
memory_classes = {
"RandomMemory": RandomMemory,
}
class_name = memory_cfg.class_name
if class_name not in memory_classes:
raise ValueError(f"Unsupported memory class_name: {class_name}. Supported: {list(memory_classes.keys())}")
MemoryClass = memory_classes[class_name]
return MemoryClass(memory_size=memory_size, num_envs=num_envs, device=device)
def _add_runtime_config(
cfg: dict,
env: Wrapper,
log_dir: str = None,
) -> dict:
# configure and instantiate the agent (visit its documentation to see all the options)
# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#configuration-and-hyperparameters
cfg = PPO_DEFAULT_CONFIG.copy()
cfg["rollouts"] = rlcfg.rollouts # memory_size
cfg["learning_epochs"] = rlcfg.learning_epochs
cfg["mini_batches"] = rlcfg.mini_batches # mini_batch_size = rollouts * num_envs / mini_batches
cfg["discount_factor"] = rlcfg.discount_factor
cfg["lambda"] = rlcfg.lambda_param
cfg["learning_rate"] = rlcfg.learning_rate
cfg["learning_rate_scheduler"] = KLAdaptiveRL
cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": rlcfg.learning_rate_scheduler_kl_threshold}
cfg["random_timesteps"] = rlcfg.random_timesteps
cfg["learning_starts"] = rlcfg.learning_starts
cfg["grad_norm_clip"] = rlcfg.grad_norm_clip
cfg["ratio_clip"] = rlcfg.ratio_clip
cfg["value_clip"] = rlcfg.value_clip
cfg["clip_predicted_values"] = rlcfg.clip_predicted_values
cfg["entropy_loss_scale"] = rlcfg.entropy_loss_scale
cfg["value_loss_scale"] = rlcfg.value_loss_scale
cfg["kl_threshold"] = rlcfg.kl_threshold
if rlcfg.rewards_shaper_scale != 1.0:
cfg["rewards_shaper"] = lambda reward, timestep, timesteps: reward * rlcfg.rewards_shaper_scale
"""Add runtime-specific configuration to the base agent config.
Args:
cfg: Base configuration from agent.to_dict() (will be modified in-place)
env: SKRL environment wrapper
log_dir: Optional logging directory path
Returns:
The same cfg dict with runtime values added (modified in-place for convenience)
"""
# Convert learning_rate_scheduler from string to actual class (if configured)
if cfg.get("learning_rate_scheduler") == "KLAdaptiveLR":
cfg["learning_rate_scheduler"] = KLAdaptiveRL
# Otherwise keep as-is (None or other scheduler type)
# Add rewards shaper (conditional based on rewards_shaper_scale in cfg)
if cfg.get("rewards_shaper_scale", 1.0) != 1.0:
cfg["rewards_shaper"] = lambda reward, timestep, timesteps: reward * cfg["rewards_shaper_scale"]
else:
cfg["rewards_shaper"] = None
cfg["time_limit_bootstrap"] = rlcfg.time_limit_bootstrap
# Add preprocessors (require runtime env values)
cfg["state_preprocessor"] = RunningStandardScaler
cfg["state_preprocessor_kwargs"] = {
"size": env.observation_space,
@@ -72,10 +97,15 @@ def _get_cfg(
}
cfg["value_preprocessor"] = RunningStandardScaler
cfg["value_preprocessor_kwargs"] = {"size": 1, "device": env.device}
# logging to TensorBoard and write checkpoints (in timesteps)
# Add experiment configuration (handle -1 -> "auto" conversion)
if log_dir:
cfg["experiment"]["write_interval"] = rlcfg.check_point_interval
cfg["experiment"]["checkpoint_interval"] = rlcfg.check_point_interval
cfg["experiment"]["write_interval"] = (
"auto" if cfg["experiment"]["write_interval"] == -1 else cfg["experiment"]["write_interval"]
)
cfg["experiment"]["checkpoint_interval"] = (
"auto" if cfg["experiment"]["checkpoint_interval"] == -1 else cfg["experiment"]["checkpoint_interval"]
)
cfg["experiment"]["directory"] = log_dir
else:
cfg["experiment"]["write_interval"] = 0
@@ -146,7 +176,7 @@ class Trainer:
_trainer: SequentialTrainer
_env_name: str
_sim_backend: str
_rlcfg: PPOCfg
_rlcfg: SkrlCfg
_enable_render: bool
def __init__(
@@ -158,7 +188,7 @@ class Trainer:
) -> None:
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
if cfg_override is not None:
rlcfg = rlcfg.replace(**cfg_override)
rlcfg = utils.cfg_override(rlcfg, cfg_override)
self._rlcfg = rlcfg
self._env_name = env_name
self._sim_backend = sim_backend
@@ -171,13 +201,16 @@ class Trainer:
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
set_seed(rlcfg.seed)
set_seed(rlcfg.runner.seed)
skrl_env = wrap_env(env, self._enable_render)
models = self._make_model(skrl_env, rlcfg)
ppo_cfg = _get_cfg(rlcfg, skrl_env, log_dir=get_log_dir(self._env_name))
agent = self._make_agent(models, skrl_env, ppo_cfg)
# Get base configuration from config object
ppo_cfg = rlcfg.runner.agent.to_dict()
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
cfg_trainer = {
"timesteps": rlcfg.max_batch_env_steps,
"timesteps": rlcfg.runner.trainer.timesteps,
"headless": not self._enable_render,
}
trainer = SequentialTrainer(cfg=cfg_trainer, env=skrl_env, agents=agent)
@@ -189,11 +222,14 @@ class Trainer:
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
set_seed(rlcfg.seed)
set_seed(rlcfg.runner.seed)
env = wrap_env(env, self._enable_render)
models = self._make_model(env, rlcfg)
ppo_cfg = _get_cfg(rlcfg, env)
agent = self._make_agent(models, env, ppo_cfg)
# Get base configuration from config object
ppo_cfg = rlcfg.runner.agent.to_dict()
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, env)
agent = self._make_agent(models, env, ppo_cfg, rlcfg.runner.memory)
agent.load(policy)
obs, _ = env.reset()
@@ -208,57 +244,63 @@ class Trainer:
if delta_time < 1.0 / fps:
time.sleep(1.0 / fps - delta_time)
def _make_model(self, env: Wrapper, rlcfg: PPOCfg) -> dict[str, Model]:
# define models (stochastic and deterministic models) using mixins
def _make_model(self, env: Wrapper, rlcfg: SkrlCfg) -> dict[str, Model]:
_activation_fn = {
"elu": nn.elu,
"relu": nn.relu,
"tanh": nn.tanh,
"sigmoid": nn.sigmoid,
"leaky_relu": nn.leaky_relu,
"selu": nn.selu,
}
policy_cfg = rlcfg.runner.models.policy
value_cfg = rlcfg.runner.models.value
def resolve_activations(activation_names: list[str], hiddens: list[int]) -> list:
if len(activation_names) == 1:
return [_activation_fn[activation_names[0]]] * len(hiddens)
if len(activation_names) != len(hiddens):
raise ValueError(
f"hidden_activation length ({len(activation_names)}) must be 1 or "
f"match hiddens length ({len(hiddens)})"
)
return [_activation_fn[name] for name in activation_names]
policy_acts = resolve_activations(policy_cfg.hidden_activation, policy_cfg.hiddens)
value_acts = resolve_activations(value_cfg.hidden_activation, value_cfg.hiddens)
class Policy(GaussianMixin, Model):
def __init__(
self,
observation_space,
action_space,
device=None,
clip_actions=False,
clip_log_std=True,
min_log_std=-20,
max_log_std=2,
reduction="sum",
**kwargs,
):
def __init__(self, observation_space, action_space, device=None, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
GaussianMixin.__init__(
self,
clip_actions,
clip_log_std,
min_log_std,
max_log_std,
reduction,
policy_cfg.clip_actions,
policy_cfg.clip_log_std,
policy_cfg.min_log_std,
policy_cfg.max_log_std,
policy_cfg.reduction,
)
@nn.compact # marks the given module method allowing inlined submodules
@nn.compact
def __call__(self, inputs, role):
x = inputs["states"]
for size in rlcfg.policy_hidden_layer_sizes:
x = nn.elu(nn.Dense(size)(x))
for size, act in zip(policy_cfg.hiddens, policy_acts):
x = act(nn.Dense(size)(x))
x = nn.Dense(self.num_actions)(x)
log_std = self.param("log_std", lambda _: jnp.ones(self.num_actions))
log_std = self.param("log_std", lambda _: jnp.full(self.num_actions, float(policy_cfg.initial_log_std)))
return x, log_std, {}
class Value(DeterministicMixin, Model):
def __init__(
self,
observation_space,
action_space,
device=None,
clip_actions=False,
**kwargs,
):
def __init__(self, observation_space, action_space, device=None, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
DeterministicMixin.__init__(self, clip_actions)
DeterministicMixin.__init__(self, value_cfg.clip_actions)
@nn.compact # marks the given module method allowing inlined submodules
@nn.compact
def __call__(self, inputs, role):
x = inputs["states"]
for size in rlcfg.value_hidden_layer_sizes:
x = nn.elu(nn.Dense(size)(x))
for size, act in zip(value_cfg.hiddens, value_acts):
x = act(nn.Dense(size)(x))
x = nn.Dense(1)(x)
return x, {}
@@ -268,21 +310,26 @@ class Trainer:
action_space=env.action_space,
device=env.device,
)
models["value"] = Value(
observation_space=env.observation_space,
action_space=env.action_space,
device=env.device,
)
# instantiate models' state dict
for role, model in models.items():
model.init_state_dict(role)
return models
def _make_agent(self, models: dict[str, Model], env: Wrapper, ppo_cfg: dict[str, Any]) -> PPO:
memory = RandomMemory(memory_size=ppo_cfg["rollouts"], num_envs=env.num_envs, device=env.device)
def _make_agent(
self, models: dict[str, Model], env: Wrapper, ppo_cfg: dict[str, Any], memory_cfg: SkrlMemoryCfg
) -> PPO:
# Use memory_size from SkrlMemoryCfg, fall back to rollouts if -1
memory_size = memory_cfg.memory_size
if memory_size == -1:
memory_size = ppo_cfg["rollouts"]
memory = _instantiate_memory(memory_cfg, memory_size, env.num_envs, env.device)
agent = PPO(
models=models,

366
motrix_rl/src/motrix_rl/skrl/torch/train/ppo.py
View File

@@ -18,9 +18,7 @@ from typing import Any
import torch
import torch.nn as nn
from skrl.agents.torch.ppo import PPO as BasePPO
from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
from skrl.envs.torch import Wrapper
from skrl.memories.torch import RandomMemory
from skrl.models.torch import DeterministicMixin, GaussianMixin, Model
from skrl.resources.preprocessors.torch import RunningStandardScaler
from skrl.resources.schedulers.torch import KLAdaptiveRL
@@ -28,42 +26,69 @@ from skrl.trainers.torch import SequentialTrainer
from skrl.utils import set_seed
from motrix_envs import registry as env_registry
from motrix_rl import registry
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir
from motrix_rl.skrl.cfg import PPOCfg
from motrix_rl.skrl.config import SkrlCfg, SkrlMemoryCfg
from motrix_rl.skrl.torch import wrap_env
def _get_cfg(
rlcfg: PPOCfg,
def _instantiate_memory(memory_cfg: SkrlMemoryCfg, memory_size: int, num_envs: int, device) -> Any:
"""Instantiate a SKRL Memory class based on configuration.
Args:
memory_cfg: Memory configuration with class_name and settings
memory_size: Size of the memory buffer
num_envs: Number of parallel environments
device: Device to place memory on
Returns:
Instantiated SKRL Memory object
Raises:
ValueError: If class_name is not supported
"""
from skrl.memories.torch import RandomMemory
# Map class_name to actual Memory class
memory_classes = {
"RandomMemory": RandomMemory,
}
class_name = memory_cfg.class_name
if class_name not in memory_classes:
raise ValueError(f"Unsupported memory class_name: {class_name}. Supported: {list(memory_classes.keys())}")
MemoryClass = memory_classes[class_name]
return MemoryClass(memory_size=memory_size, num_envs=num_envs, device=device)
def _add_runtime_config(
cfg: dict,
env: Wrapper,
log_dir: str = None,
) -> dict:
# configure and instantiate the agent (visit its documentation to see all the options)
# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#configuration-and-hyperparameters
cfg = PPO_DEFAULT_CONFIG.copy()
cfg["rollouts"] = rlcfg.rollouts # memory_size
cfg["learning_epochs"] = rlcfg.learning_epochs
cfg["mini_batches"] = rlcfg.mini_batches # mini_batch_size = rollouts * num_envs / mini_batches
cfg["discount_factor"] = rlcfg.discount_factor
cfg["lambda"] = rlcfg.lambda_param
cfg["learning_rate"] = rlcfg.learning_rate
cfg["learning_rate_scheduler"] = KLAdaptiveRL
cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": rlcfg.learning_rate_scheduler_kl_threshold}
cfg["random_timesteps"] = rlcfg.random_timesteps
cfg["learning_starts"] = rlcfg.learning_starts
cfg["grad_norm_clip"] = rlcfg.grad_norm_clip
cfg["ratio_clip"] = rlcfg.ratio_clip
cfg["value_clip"] = rlcfg.value_clip
cfg["clip_predicted_values"] = rlcfg.clip_predicted_values
cfg["entropy_loss_scale"] = rlcfg.entropy_loss_scale
cfg["value_loss_scale"] = rlcfg.value_loss_scale
cfg["kl_threshold"] = rlcfg.kl_threshold
if rlcfg.rewards_shaper_scale != 1.0:
cfg["rewards_shaper"] = lambda reward, timestep, timesteps: reward * rlcfg.rewards_shaper_scale
"""Add runtime-specific configuration to the base agent config.
Args:
cfg: Base configuration from agent.to_dict() (will be modified in-place)
env: SKRL environment wrapper
log_dir: Optional logging directory path
Returns:
The same cfg dict with runtime values added (modified in-place for convenience)
"""
# Convert learning_rate_scheduler from string to actual class (if configured)
if cfg.get("learning_rate_scheduler") == "KLAdaptiveLR":
cfg["learning_rate_scheduler"] = KLAdaptiveRL
# Otherwise keep as-is (None or other scheduler type)
# Add rewards shaper (conditional based on rewards_shaper_scale in cfg)
if cfg.get("rewards_shaper_scale", 1.0) != 1.0:
cfg["rewards_shaper"] = lambda reward, timestep, timesteps: reward * cfg["rewards_shaper_scale"]
else:
cfg["rewards_shaper"] = None
cfg["time_limit_bootstrap"] = rlcfg.time_limit_bootstrap
# Add preprocessors (require runtime env values)
cfg["state_preprocessor"] = RunningStandardScaler
cfg["state_preprocessor_kwargs"] = {
"size": env.observation_space,
@@ -71,10 +96,15 @@ def _get_cfg(
}
cfg["value_preprocessor"] = RunningStandardScaler
cfg["value_preprocessor_kwargs"] = {"size": 1, "device": env.device}
# logging to TensorBoard and write checkpoints (in timesteps)
# Add experiment configuration (handle -1 -> "auto" conversion)
if log_dir:
cfg["experiment"]["write_interval"] = rlcfg.check_point_interval
cfg["experiment"]["checkpoint_interval"] = rlcfg.check_point_interval
cfg["experiment"]["write_interval"] = (
"auto" if cfg["experiment"]["write_interval"] == -1 else cfg["experiment"]["write_interval"]
)
cfg["experiment"]["checkpoint_interval"] = (
"auto" if cfg["experiment"]["checkpoint_interval"] == -1 else cfg["experiment"]["checkpoint_interval"]
)
cfg["experiment"]["directory"] = log_dir
else:
cfg["experiment"]["write_interval"] = 0
@@ -146,7 +176,7 @@ class Trainer:
_trainer: SequentialTrainer
_env_name: str
_sim_backend: str
_rlcfg: PPOCfg
_rlcfg: SkrlCfg
_enable_render: bool
def __init__(
@@ -158,7 +188,7 @@ class Trainer:
) -> None:
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="torch")
if cfg_override is not None:
rlcfg = rlcfg.replace(**cfg_override)
rlcfg = utils.cfg_override(rlcfg, cfg_override)
self._rlcfg = rlcfg
self._env_name = env_name
self._sim_backend = sim_backend
@@ -170,13 +200,16 @@ class Trainer:
"""
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
set_seed(rlcfg.seed)
set_seed(rlcfg.runner.seed)
skrl_env = wrap_env(env, self._enable_render)
models = self._make_model(skrl_env, rlcfg)
ppo_cfg = _get_cfg(rlcfg, skrl_env, log_dir=get_log_dir(self._env_name))
agent = self._make_agent(models, skrl_env, ppo_cfg)
# Get base configuration from config object
ppo_cfg = rlcfg.runner.agent.to_dict()
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
cfg_trainer = {
"timesteps": rlcfg.max_batch_env_steps,
"timesteps": rlcfg.runner.trainer.timesteps,
"headless": not self._enable_render,
}
trainer = SequentialTrainer(cfg=cfg_trainer, env=skrl_env, agents=agent)
@@ -187,11 +220,14 @@ class Trainer:
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
set_seed(rlcfg.seed)
set_seed(rlcfg.runner.seed)
env = wrap_env(env, self._enable_render)
models = self._make_model(env, rlcfg)
ppo_cfg = _get_cfg(rlcfg, env)
agent = self._make_agent(models, env, ppo_cfg)
# Get base configuration from config object
ppo_cfg = rlcfg.runner.agent.to_dict()
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, env)
agent = self._make_agent(models, env, ppo_cfg, rlcfg.runner.memory)
agent.load(policy)
with torch.no_grad():
obs, _ = env.reset()
@@ -206,143 +242,139 @@ class Trainer:
if delta_time < 1.0 / fps:
time.sleep(1.0 / fps - delta_time)
def _make_model(self, env: Wrapper, rlcfg: PPOCfg) -> dict[str, Model]:
def build_mlp(
input_size: int,
hidden_sizes: tuple[int, ...],
output_size: int,
activation=nn.ELU,
):
"""Helper function to build MLP layers."""
def _make_model(self, env: Wrapper, rlcfg: SkrlCfg) -> dict[str, Model]:
_activation_fn = {
"elu": nn.ELU,
"relu": nn.ReLU,
"tanh": nn.Tanh,
"sigmoid": nn.Sigmoid,
"leaky_relu": nn.LeakyReLU,
"selu": nn.SELU,
}
policy_cfg = rlcfg.runner.models.policy
value_cfg = rlcfg.runner.models.value
separate = rlcfg.runner.models.separate
def resolve_activations(activation_names: list[str], hiddens: list[int]) -> list:
if len(activation_names) == 1:
return [_activation_fn[activation_names[0]]] * len(hiddens)
if len(activation_names) != len(hiddens):
raise ValueError(
f"hidden_activation length ({len(activation_names)}) must be 1 or "
f"match hiddens length ({len(hiddens)})"
)
return [_activation_fn[name] for name in activation_names]
policy_acts = resolve_activations(policy_cfg.hidden_activation, policy_cfg.hiddens)
value_acts = resolve_activations(value_cfg.hidden_activation, value_cfg.hiddens)
def build_mlp(input_size: int, hidden_sizes: list[int], activations: list) -> nn.Sequential:
layers = []
current_size = input_size
for hidden_size in hidden_sizes:
for hidden_size, act in zip(hidden_sizes, activations):
layers.append(nn.Linear(current_size, hidden_size))
layers.append(activation())
layers.append(act())
current_size = hidden_size
layers.append(nn.Linear(current_size, output_size))
return nn.Sequential(*layers)
# define shared model (stochastic and deterministic models) using mixins
class Shared(GaussianMixin, DeterministicMixin, Model):
def __init__(
self,
observation_space,
action_space,
device,
policy_hidden_sizes,
value_hidden_sizes,
share_features=True,
clip_actions=False,
clip_log_std=True,
min_log_std=-20,
max_log_std=2,
reduction="sum",
):
Model.__init__(self, observation_space, action_space, device)
GaussianMixin.__init__(
self,
clip_actions,
clip_log_std,
min_log_std,
max_log_std,
reduction,
)
DeterministicMixin.__init__(self, clip_actions)
# Use configured share_features setting
self.share_features = share_features and policy_hidden_sizes == value_hidden_sizes
if self.share_features:
# Build shared feature extraction layers
shared_layers = []
current_size = self.num_observations
for hidden_size in policy_hidden_sizes:
shared_layers.append(nn.Linear(current_size, hidden_size))
shared_layers.append(nn.ELU())
current_size = hidden_size
self.net = nn.Sequential(*shared_layers)
self.mean_layer = nn.Linear(current_size, self.num_actions)
self.log_std_parameter = nn.Parameter(torch.ones(self.num_actions))
self.value_layer = nn.Linear(current_size, 1)
else:
# Build separate networks for policy and value
self.policy_net = build_mlp(
self.num_observations,
policy_hidden_sizes[:-1],
policy_hidden_sizes[-1] if len(policy_hidden_sizes) > 0 else self.num_actions,
)
self.value_net = build_mlp(
self.num_observations,
value_hidden_sizes[:-1],
value_hidden_sizes[-1] if len(value_hidden_sizes) > 0 else 1,
)
# Output layers
if len(policy_hidden_sizes) > 0:
self.mean_layer = nn.Linear(policy_hidden_sizes[-1], self.num_actions)
else:
self.mean_layer = nn.Linear(self.num_observations, self.num_actions)
self.log_std_parameter = nn.Parameter(torch.ones(self.num_actions))
if len(value_hidden_sizes) > 0:
self.value_layer = nn.Linear(value_hidden_sizes[-1], 1)
else:
self.value_layer = nn.Linear(self.num_observations, 1)
def act(self, inputs, role):
if role == "policy":
return GaussianMixin.act(self, inputs, role)
elif role == "value":
return DeterministicMixin.act(self, inputs, role)
def compute(self, inputs, role):
if role == "policy":
if self.share_features:
self._shared_output = self.net(inputs["states"])
return (
self.mean_layer(self._shared_output),
self.log_std_parameter,
{},
)
else:
policy_features = self.policy_net(inputs["states"])
return (
self.mean_layer(policy_features),
self.log_std_parameter,
{},
)
elif role == "value":
if self.share_features:
shared_output = (
self.net(inputs["states"]) if self._shared_output is None else self._shared_output
)
self._shared_output = None
return self.value_layer(shared_output), {}
else:
value_features = self.value_net(inputs["states"])
return self.value_layer(value_features), {}
models = {}
models["policy"] = Shared(
observation_space=env.observation_space,
action_space=env.action_space,
device=env.device,
policy_hidden_sizes=rlcfg.policy_hidden_layer_sizes,
value_hidden_sizes=rlcfg.value_hidden_layer_sizes,
share_features=rlcfg.share_policy_value_features,
)
models["value"] = models["policy"]
if separate:
class Policy(GaussianMixin, Model):
def __init__(self, observation_space, action_space, device, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
GaussianMixin.__init__(
self,
policy_cfg.clip_actions,
policy_cfg.clip_log_std,
policy_cfg.min_log_std,
policy_cfg.max_log_std,
policy_cfg.reduction,
)
self.net = build_mlp(self.num_observations, policy_cfg.hiddens, policy_acts)
self.mean_layer = nn.Linear(policy_cfg.hiddens[-1], self.num_actions)
self.log_std_parameter = nn.Parameter(torch.full((self.num_actions,), policy_cfg.initial_log_std))
def compute(self, inputs, role):
x = self.net(inputs["states"])
return self.mean_layer(x), self.log_std_parameter, {}
class Value(DeterministicMixin, Model):
def __init__(self, observation_space, action_space, device, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
DeterministicMixin.__init__(self, value_cfg.clip_actions)
self.net = build_mlp(self.num_observations, value_cfg.hiddens, value_acts)
self.value_layer = nn.Linear(value_cfg.hiddens[-1], 1)
def compute(self, inputs, role):
x = self.net(inputs["states"])
return self.value_layer(x), {}
models["policy"] = Policy(
observation_space=env.observation_space,
action_space=env.action_space,
device=env.device,
)
models["value"] = Value(
observation_space=env.observation_space,
action_space=env.action_space,
device=env.device,
)
else:
class Shared(GaussianMixin, DeterministicMixin, Model):
def __init__(self, observation_space, action_space, device, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
GaussianMixin.__init__(
self,
policy_cfg.clip_actions,
policy_cfg.clip_log_std,
policy_cfg.min_log_std,
policy_cfg.max_log_std,
policy_cfg.reduction,
)
DeterministicMixin.__init__(self, value_cfg.clip_actions)
self.net = build_mlp(self.num_observations, policy_cfg.hiddens, policy_acts)
self.mean_layer = nn.Linear(policy_cfg.hiddens[-1], self.num_actions)
self.log_std_parameter = nn.Parameter(torch.full((self.num_actions,), policy_cfg.initial_log_std))
self.value_layer = nn.Linear(policy_cfg.hiddens[-1], 1)
self._shared_output = None
def act(self, inputs, role):
if role == "policy":
return GaussianMixin.act(self, inputs, role)
elif role == "value":
return DeterministicMixin.act(self, inputs, role)
def compute(self, inputs, role):
if role == "policy":
self._shared_output = self.net(inputs["states"])
return self.mean_layer(self._shared_output), self.log_std_parameter, {}
elif role == "value":
shared = self._shared_output if self._shared_output is not None else self.net(inputs["states"])
self._shared_output = None
return self.value_layer(shared), {}
models["policy"] = Shared(
observation_space=env.observation_space,
action_space=env.action_space,
device=env.device,
)
models["value"] = models["policy"]
return models
def _make_agent(self, models: dict[str, Model], env: Wrapper, ppo_cfg: dict[str, Any]) -> PPO:
memory = RandomMemory(memory_size=ppo_cfg["rollouts"], num_envs=env.num_envs, device=env.device)
def _make_agent(
self, models: dict[str, Model], env: Wrapper, ppo_cfg: dict[str, Any], memory_cfg: SkrlMemoryCfg
) -> PPO:
# Use memory_size from SkrlMemoryCfg, fall back to rollouts if -1
memory_size = memory_cfg.memory_size
if memory_size == -1:
memory_size = ppo_cfg["rollouts"]
memory = _instantiate_memory(memory_cfg, memory_size, env.num_envs, env.device)
agent = PPO(
models=models,

38
motrix_rl/src/motrix_rl/tasks/__init__.py Normal file
View File

@@ -0,0 +1,38 @@
# 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.
# ==============================================================================
from . import (
acrobot, # noqa: F401
anymal_navigation, # noqa: F401
bounce_ball, # noqa: F401
cartpole, # noqa: F401
dm_cheetah, # noqa: F401
dm_finger, # noqa: F401
dm_hopper, # noqa: F401
dm_humanoid, # noqa: F401
dm_lqr, # noqa: F401
dm_manipulator, # noqa: F401
dm_quadruped, # noqa: F401
dm_reacher, # noqa: F401
dm_walker, # noqa: F401
franka_lift_cube, # noqa: F401
franka_open_cabinet, # noqa: F401
go1, # noqa: F401
go2, # noqa: F401
pendulum, # noqa: F401
point_mass, # noqa: F401
rm65_open_cabinet, # noqa: F401
shadow_hand_repose, # noqa: F401
)

123
motrix_rl/src/motrix_rl/tasks/acrobot.py Normal file
View File

@@ -0,0 +1,123 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
# -- docs-tag-start: acrobot-train-cfg --
@rlcfg("acrobot", backend="jax")
@dataclass
class AcrobotSkrlPpo(SkrlCfg):
"""Acrobot SKRL configuration with nested structure.
Configuration overrides:
- Network architecture: 32x32 hidden layers for both policy and value
- PPO parameters: 64 rollouts, 5 learning epochs, 8 mini-batches
- Training: 60M timesteps
"""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
# Configure model architectures
runner.models.policy.hiddens = [32, 32]
runner.models.value.hiddens = [32, 32]
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 64
agent.learning_epochs = 5
agent.mini_batches = 8
agent.learning_rate = 3e-4
agent.grad_norm_clip = 0.1
agent.entropy_loss_scale = 0.1
agent.discount_factor = 0.995
agent.lam = 0.97
agent.ratio_clip = 0.2
agent.value_loss_scale = 0.5
agent.value_clip = 10.0
agent.clip_predicted_values = False
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.02}
agent.kl_threshold = 0.03
# Configure training parameters
runner.trainer.timesteps = 29000
# -- docs-tag-end: acrobot-train-cfg --
@rlcfg("acrobot", backend="torch")
@dataclass
class AcrobotSkrlPpoTorch(SkrlCfg):
"""Acrobot SKRL Torch configuration with nested structure.
Configuration overrides:
- Network architecture: 32x32 hidden layers for both policy and value
- PPO parameters: 64 rollouts, 5 learning epochs, 8 mini-batches
- Training: 60M timesteps
- Torch-specific: entropy_loss_scale=0.2 (vs 0.1 for JAX)
"""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
# Configure model architectures
runner.models.policy.hiddens = [32, 32]
runner.models.value.hiddens = [32, 32]
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 64
agent.learning_epochs = 5
agent.mini_batches = 8
agent.learning_rate = 3e-4
agent.grad_norm_clip = 0.1
agent.entropy_loss_scale = 0.2 # Torch-specific: higher than JAX (0.1)
agent.discount_factor = 0.995
agent.lam = 0.97
agent.ratio_clip = 0.2
agent.value_loss_scale = 0.5
agent.value_clip = 10.0
agent.clip_predicted_values = False
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.02}
agent.kl_threshold = 0.03
# Configure training parameters
runner.trainer.timesteps = 29000
@rlcfg("acrobot")
@dataclass
class AcrobotRslrlPpo(RslrlCfg):
"""Acrobot RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 900
runner.num_steps_per_env = 32
runner.experiment_name = "acrobot"
runner.actor.hidden_dims = [32, 32]
runner.critic.hidden_dims = [32, 32]
algo.learning_rate = 1e-4
algo.entropy_coef = 0.005

95
motrix_rl/src/motrix_rl/tasks/anymal_navigation.py Normal file
View File

@@ -0,0 +1,95 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("anymal_c_navigation_flat")
@dataclass
class AnymalCPPO(SkrlCfg):
"""Anymal C Navigation SKRL PPO configuration.
Configuration for training ANYmal C robot for flat terrain navigation.
Uses medium-sized network suitable for most locomotion tasks.
"""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
self.num_envs = 2048
self.play_num_envs = 16
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# ===== Basic Training Parameters =====
runner.seed = 42
# ===== Network Architecture =====
# Medium-sized network (default configuration, suitable for most tasks)
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# ===== PPO Core Parameters =====
agent.rollouts = 48
agent.learning_epochs = 6
agent.mini_batches = 32
agent.learning_rate = 3e-4
agent.discount_factor = 0.99
agent.lam = 0.95
agent.grad_norm_clip = 1.0
# ===== PPO Clipping Parameters =====
agent.ratio_clip = 0.2
agent.value_clip = 0.2
agent.clip_predicted_values = True
# ===== Training Parameters =====
trainer.timesteps = 48000
class rslrl:
@rlcfg("anymal_c_navigation_flat")
@dataclass
class AnymalCPpoRslrl(RslrlCfg):
"""Anymal C Navigation RSLRL PPO configuration."""
def __post_init__(self):
"""Configure RSLRL runner and algorithm settings."""
self.num_envs = 2048
self.play_num_envs = 16
runner = self.runner
algo = runner.algorithm
# ===== Basic Training Parameters =====
runner.seed = 42
# max_iterations = max_env_steps / num_envs / roll_out = 100000000 / 2048 / 48 ≈ 1017
runner.max_iterations = 1017
runner.num_steps_per_env = 48
runner.experiment_name = "anymal_c_navigation_flat"
# ===== Network Architecture =====
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
# ===== Algorithm Parameters =====
algo.learning_rate = 3e-4
algo.num_learning_epochs = 6
algo.num_mini_batches = 4

66
motrix_rl/src/motrix_rl/tasks/bounce_ball.py Normal file
View File

@@ -0,0 +1,66 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("bounce_ball")
@dataclass
class BounceBallSkrlPpo(SkrlCfg):
"""BounceBall SKRL configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
runner.seed = 42
# Configure model architectures
runner.models.policy.hiddens = [64, 64, 64]
runner.models.value.hiddens = [64, 64, 64]
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 128
agent.learning_epochs = 15
agent.mini_batches = 16
agent.learning_rate = 2e-4
# Configure training parameters
# trainer.timesteps = max_env_steps / num_envs
runner.trainer.timesteps = 45000
@rlcfg("bounce_ball")
@dataclass
class BounceBallRslrlPpo(RslrlCfg):
"""BounceBall RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 400
runner.num_steps_per_env = 128
runner.experiment_name = "bounce_ball"
runner.actor.hidden_dims = [64, 64, 64]
runner.critic.hidden_dims = [64, 64, 64]
algo.learning_rate = 1e-3
algo.entropy_coef = 0.005
algo.num_learning_epochs = 2
algo.num_mini_batches = 4

189
motrix_rl/src/motrix_rl/tasks/cartpole.py Normal file
View File

@@ -0,0 +1,189 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import (
RslrlCfg,
)
from motrix_rl.skrl.config import (
SkrlCfg,
)
# docs-start: cartpole-skrl-config
@rlcfg("cartpole")
@dataclass
class CartPoleSkrlPpo(SkrlCfg):
"""CartPole SKRL configuration with complete explicit parameter filling.
All parameters from parent classes are explicitly specified.
"""
def __post_init__(self):
"""Configure SKRL runner settings with explicit parameters."""
# Environment settings
self.num_envs = 2048 # Number of parallel environments during training
self.play_num_envs = 16 # Number of parallel environments during evaluation
# Get runner and nested configs
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Random seed
runner.seed = 42 # Random seed for reproducibility
# Models configuration
models.separate = False # Share features between policy and value networks
# Policy network configuration
models.policy.class_name = "GaussianMixin" # Use Gaussian policy for continuous actions
models.policy.clip_actions = False # Don't clip actions to action space
models.policy.clip_log_std = True # Clip log standard deviation
models.policy.initial_log_std = 1.0 # Initial log standard deviation
models.policy.min_log_std = -20.0 # Minimum log standard deviation
models.policy.max_log_std = 2.0 # Maximum log standard deviation
models.policy.reduction = "sum" # Reduction method for loss computation
models.policy.input = "STATES" # Input to policy network
models.policy.hiddens = [32, 32] # Hidden layer sizes (small network for simple task)
models.policy.hidden_activation = ["elu"] # Activation function for hidden layers
models.policy.output = "ACTIONS" # Output of policy network
models.policy.output_activation = "" # No activation for output layer
models.policy.output_scale = 1.0 # Scale factor for output
# Value network configuration
models.value.class_name = "DeterministicMixin" # Use deterministic value function
models.value.clip_actions = False # Don't clip actions
models.value.input = "STATES" # Input to value network
models.value.hiddens = [32, 32] # Hidden layer sizes (small network for simple task)
models.value.hidden_activation = ["elu"] # Activation function for hidden layers
models.value.output = "ONE" # Output single value
models.value.output_activation = "" # No activation for output layer
models.value.output_scale = 1.0 # Scale factor for output
# Memory configuration
runner.memory.class_name = "RandomMemory" # Use random sampling memory
runner.memory.memory_size = -1 # Unlimited memory size (-1 means auto-calculate)
# Agent configuration
agent.class_name = "PPO" # Use Proximal Policy Optimization algorithm
agent.rollouts = 32 # Number of experience rollouts to collect
agent.learning_epochs = 5 # Number of learning epochs per update (higher than default 2)
agent.mini_batches = 4 # Number of mini-batches (fewer than default 32 for simple task)
agent.discount_factor = 0.99 # Discount factor (gamma) for future rewards
agent.lam = 0.95 # GAE (Generalized Advantage Estimation) lambda parameter
agent.learning_rate = 1e-3 # Learning rate for optimizer
agent.learning_rate_scheduler = "KLAdaptiveLR" # Use KL-divergence adaptive learning rate
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.008} # KL threshold for adaptive LR
agent.random_timesteps = 0 # Number of random timesteps before using policy
agent.learning_starts = 0 # Timesteps before learning starts
agent.grad_norm_clip = 1.0 # Maximum gradient norm for clipping
agent.ratio_clip = 0.2 # PPO clipping ratio for policy update
agent.value_clip = 0.2 # Clipping parameter for value function loss
agent.clip_predicted_values = True # Clip predicted values in value loss
agent.entropy_loss_scale = 0.0 # Coefficient for entropy loss (disabled)
agent.value_loss_scale = 2.0 # Coefficient for value function loss
agent.kl_threshold = 0 # KL divergence threshold (0 means disabled)
agent.rewards_shaper_scale = 1.0 # Scale factor for reward shaping
agent.time_limit_bootstrap = True # Use bootstrapping for time-limited episodes
# Experiment configuration
agent.experiment.directory = "runs" # Directory to save experiment results
agent.experiment.experiment_name = "" # Experiment name (empty means auto-generated)
agent.experiment.write_interval = -1 # TensorBoard write interval (-1 means default)
agent.experiment.checkpoint_interval = -1 # Checkpoint save interval (-1 means default)
# Trainer configuration
trainer.class_name = "SequentialTrainer" # Use sequential trainer
trainer.timesteps = 5000 # Total training timesteps (sufficient for CartPole)
# docs-end: cartpole-skrl-config
# docs-start: cartpole-rslrl-config
@rlcfg("cartpole")
@dataclass
class CartPoleRslrlPpo(RslrlCfg):
"""CartPole RSLRL configuration with complete explicit parameter filling.
All parameters from parent classes are explicitly specified.
"""
def __post_init__(self):
"""Configure RSLRL runner settings with explicit parameters."""
# Environment settings
self.num_envs = 2048 # Number of parallel environments during training
self.play_num_envs = 16 # Number of parallel environments during evaluation
# Get runner and nested configs
runner = self.runner
actor = runner.actor
critic = runner.critic
algo = runner.algorithm
# Runner settings
runner.class_name = "OnPolicyRunner" # Use on-policy runner
runner.seed = 42 # Random seed for reproducibility
runner.device = "cuda:0" # Device to use for training
runner.num_steps_per_env = 16 # Number of steps to collect per environment
runner.max_iterations = 300 # Total number of training iterations
runner.save_interval = 50 # Checkpoint save interval
runner.experiment_name = "cartpole" # Experiment name for logging
runner.run_name = "" # Run name (empty means auto-generated)
runner.logger = "tensorboard" # Logger type
runner.obs_groups = {"actor": ["policy"], "critic": ["policy"]} # Observation groups
# Actor network configuration
actor.class_name = "MLPModel" # Use MLP model
actor.hidden_dims = [32, 32] # Hidden layer sizes (small network for simple task)
actor.activation = "elu" # Activation function for hidden layers
actor.obs_normalization = True # Normalize observations
actor.stochastic = True # Use stochastic policy
actor.init_noise_std = 1.0 # Initial noise standard deviation
actor.noise_std_type = "scalar" # Noise std type (scalar or log)
actor.state_dependent_std = False # Use state-dependent std
# Critic network configuration
critic.class_name = "MLPModel" # Use MLP model
critic.hidden_dims = [32, 32] # Hidden layer sizes (small network for simple task)
critic.activation = "elu" # Activation function for hidden layers
critic.obs_normalization = True # Normalize observations
critic.stochastic = False # Use deterministic value function
# PPO algorithm configuration
algo.class_name = "PPO" # Use PPO algorithm
algo.optimizer = "adam" # Optimizer type
algo.learning_rate = 5.0e-4 # Learning rate for optimizer
algo.num_learning_epochs = 2 # Number of learning epochs per iteration
algo.num_mini_batches = 4 # Number of mini-batches for optimization
algo.schedule = "adaptive" # Learning rate schedule
algo.value_loss_coef = 1.0 # Value loss coefficient
algo.clip_param = 0.2 # PPO clipping parameter
algo.use_clipped_value_loss = True # Use clipped value loss
algo.desired_kl = 0.008 # Desired KL divergence for adaptive learning
algo.entropy_coef = 5e-3 # Entropy coefficient for exploration
algo.gamma = 0.99 # Discount factor
algo.lam = 0.95 # GAE lambda parameter
algo.max_grad_norm = 1.0 # Maximum gradient norm for clipping
algo.normalize_advantage_per_mini_batch = False # Normalize advantage per mini-batch
algo.rnd_cfg = None # RND configuration (disabled)
algo.symmetry_cfg = None # Symmetry configuration (disabled)
# docs-end: cartpole-rslrl-config

65
motrix_rl/src/motrix_rl/tasks/dm_cheetah.py Normal file
View File

@@ -0,0 +1,65 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-cheetah")
@dataclass
class CheetahSkrlPpo(SkrlCfg):
"""Cheetah SKRL configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-cheetah")
@dataclass
class CheetahRslrlPpo(RslrlCfg):
"""Cheetah RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 900
runner.num_steps_per_env = 24
runner.experiment_name = "dm_cheetah"
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 32

120
motrix_rl/src/motrix_rl/tasks/dm_finger.py Normal file
View File

@@ -0,0 +1,120 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-finger-spin")
@rlcfg("dm-finger-turn-easy")
@rlcfg("dm-finger-turn-hard")
@dataclass
class FingerSkrlPpo(SkrlCfg):
"""Finger SKRL configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
agent = runner.agent
trainer = runner.trainer
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 4
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 20000
# @rlcfg("dm-finger-spin", backend="jax")
# @dataclass
# class FingerSpinSkrlPpoJax(SkrlCfg):
# """Finger Spin SKRL configuration for JAX backend.
# More conservative PPO for stability (spin can collapse mid-training in JAX).
# """
# def __post_init__(self):
# """Configure nested SKRL runner settings."""
# runner = self.runner
# agent = runner.agent
# trainer = runner.trainer
# # Configure PPO agent parameters (conservative for stability)
# agent.rollouts = 24
# agent.learning_epochs = 1
# agent.mini_batches = 16
# agent.learning_rate = 7.5e-5
# agent.ratio_clip = 0.08
# agent.value_clip = 0.1
# agent.value_loss_scale = 0.5
# agent.grad_norm_clip = 0.25
# agent.entropy_loss_scale = 5e-4
# agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.003}
# # Configure training parameters
# trainer.timesteps = 20000
# @rlcfg("dm-finger-turn-hard", backend="jax")
# @dataclass
# class FingerTurnHardSkrlPpoJax(SkrlCfg):
# """Finger Turn Hard SKRL configuration for JAX backend.
# Extra conservative to avoid late-stage collapses.
# """
# def __post_init__(self):
# """Configure nested SKRL runner settings."""
# runner = self.runner
# agent = runner.agent
# trainer = runner.trainer
# # Configure PPO agent parameters (extra conservative)
# agent.rollouts = 24
# agent.learning_epochs = 1
# agent.mini_batches = 16
# agent.learning_rate = 5e-5
# agent.ratio_clip = 0.08
# agent.value_loss_scale = 0.5
# agent.grad_norm_clip = 0.25
# agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.004}
# # Configure training parameters
# trainer.timesteps = 20000
@rlcfg("dm-finger-spin")
@rlcfg("dm-finger-turn-easy")
@rlcfg("dm-finger-turn-hard")
@dataclass
class FingerRslrlPpo(RslrlCfg):
"""Finger RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 500
runner.num_steps_per_env = 24
runner.experiment_name = "dm_finger"
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 4

117
motrix_rl/src/motrix_rl/tasks/dm_hopper.py Normal file
View File

@@ -0,0 +1,117 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-hopper-stand", backend="jax")
@dataclass
class HopperStandSkrlJaxPpo(SkrlCfg):
"""Hopper Stand SKRL JAX configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [32, 32, 32]
models.value.hiddens = [32, 32, 32]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 4
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-hopper-stand", backend="torch")
@rlcfg("dm-hopper-hop", backend="torch")
@rlcfg("dm-hopper-hop", backend="jax")
@dataclass
class HopperSkrlTorchPpo(SkrlCfg):
"""Hopper SKRL Torch configuration (shared for stand and hop tasks)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [32, 32, 32]
models.value.hiddens = [32, 32, 32]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 5
agent.mini_batches = 32
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 20000
# ==============================================================================
# RSLRL Configurations
# ==============================================================================
@rlcfg("dm-hopper-stand")
@dataclass
class HopperStandRslrlPpo(RslrlCfg):
"""Hopper Stand RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 833
runner.num_steps_per_env = 24
runner.experiment_name = "dm_hopper_stand"
runner.actor.hidden_dims = [32, 32, 32]
runner.critic.hidden_dims = [32, 32, 32]
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 4
@rlcfg("dm-hopper-hop")
@dataclass
class HopperHopRslrlPpo(RslrlCfg):
"""Hopper Hop RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 833
runner.num_steps_per_env = 24
runner.experiment_name = "dm_hopper_hop"
runner.actor.hidden_dims = [32, 32, 32]
runner.critic.hidden_dims = [32, 32, 32]
algo.learning_rate = 2e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 32

98
motrix_rl/src/motrix_rl/tasks/dm_humanoid.py Normal file
View File

@@ -0,0 +1,98 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-humanoid-stand", backend="jax")
@rlcfg("dm-humanoid-walk", backend="jax")
@rlcfg("dm-humanoid-run", backend="jax")
@dataclass
class HumanoidSkrlPpo(SkrlCfg):
"""Humanoid SKRL configuration with nested structure (JAX)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [512, 256, 128]
models.value.hiddens = [512, 256, 128]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 8
agent.mini_batches = 2
agent.learning_rate = 3e-4
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-humanoid-stand", backend="torch")
@rlcfg("dm-humanoid-walk", backend="torch")
@rlcfg("dm-humanoid-run", backend="torch")
@dataclass
class HumanoidSkrlPpoTorch(SkrlCfg):
"""Humanoid SKRL configuration with nested structure (PyTorch)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [512, 256, 128]
models.value.hiddens = [512, 256, 128]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 8
agent.mini_batches = 2
agent.learning_rate = 3e-4
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-humanoid-stand")
@rlcfg("dm-humanoid-walk")
@rlcfg("dm-humanoid-run")
@dataclass
class HumanoidRslrlPpo(RslrlCfg):
"""Humanoid RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 400
runner.num_steps_per_env = 64
runner.experiment_name = "dm_humanoid"
runner.actor.hidden_dims = [512, 256, 128]
runner.critic.hidden_dims = [512, 256, 128]
algo.learning_rate = 3e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 4
algo.entropy_coef = 0.001

182
motrix_rl/src/motrix_rl/tasks/dm_lqr.py Normal file
View File

@@ -0,0 +1,182 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-lqr-2-1", backend="jax")
@dataclass
class Lqr21SkrlJaxPpo(SkrlCfg):
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.hiddens = [128, 128, 64]
models.value.hiddens = [128, 128, 64]
agent.rollouts = 64
agent.learning_epochs = 8
agent.mini_batches = 8
agent.learning_rate = 1.5e-4
agent.discount_factor = 0.995
agent.lam = 0.97
agent.grad_norm_clip = 0.5
agent.ratio_clip = 0.12
agent.value_clip = 0.1
agent.value_loss_scale = 1.0
agent.entropy_loss_scale = 1e-3
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.01}
trainer.timesteps = 2500
@rlcfg("dm-lqr-2-1", backend="torch")
@dataclass
class Lqr21SkrlTorchPpo(SkrlCfg):
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.hiddens = [128, 128, 64]
models.value.hiddens = [128, 128, 64]
agent.rollouts = 64
agent.learning_epochs = 8
agent.mini_batches = 8
agent.learning_rate = 1.5e-4
agent.discount_factor = 0.995
agent.lam = 0.97
agent.grad_norm_clip = 0.5
agent.ratio_clip = 0.12
agent.value_clip = 0.1
agent.value_loss_scale = 1.0
agent.entropy_loss_scale = 1e-3
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.01}
trainer.timesteps = 2500
@rlcfg("dm-lqr-6-2", backend="jax")
@dataclass
class Lqr62SkrlJaxPpo(SkrlCfg):
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
agent.rollouts = 96
agent.learning_epochs = 8
agent.mini_batches = 16
agent.learning_rate = 1.0e-4
agent.discount_factor = 0.995
agent.lam = 0.97
agent.grad_norm_clip = 0.5
agent.ratio_clip = 0.12
agent.value_clip = 0.1
agent.value_loss_scale = 1.0
agent.entropy_loss_scale = 1e-3
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.01}
trainer.timesteps = 20000
@rlcfg("dm-lqr-6-2", backend="torch")
@dataclass
class Lqr62SkrlTorchPpo(SkrlCfg):
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
agent.rollouts = 96
agent.learning_epochs = 8
agent.mini_batches = 16
agent.learning_rate = 1.0e-4
agent.discount_factor = 0.995
agent.lam = 0.97
agent.grad_norm_clip = 0.5
agent.ratio_clip = 0.12
agent.value_clip = 0.1
agent.value_loss_scale = 1.0
agent.entropy_loss_scale = 1e-3
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.01}
trainer.timesteps = 20000
@rlcfg("dm-lqr-2-1")
@dataclass
class Lqr21RslrlPpo(RslrlCfg):
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 1500
runner.num_steps_per_env = 64
runner.experiment_name = "dm_lqr_2_1"
runner.actor.hidden_dims = [128, 128, 64]
runner.critic.hidden_dims = [128, 128, 64]
algo.learning_rate = 1.5e-4
algo.num_learning_epochs = 8
algo.num_mini_batches = 8
algo.gamma = 0.995
algo.lam = 0.97
algo.clip_param = 0.12
algo.desired_kl = 0.01
algo.entropy_coef = 1e-3
algo.max_grad_norm = 0.5
@rlcfg("dm-lqr-6-2")
@dataclass
class Lqr62RslrlPpo(RslrlCfg):
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 2000
runner.num_steps_per_env = 96
runner.experiment_name = "dm_lqr_6_2"
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
algo.learning_rate = 1.0e-4
algo.num_learning_epochs = 8
algo.num_mini_batches = 16
algo.gamma = 0.995
algo.lam = 0.97
algo.clip_param = 0.12
algo.desired_kl = 0.01
algo.entropy_coef = 1e-3
algo.max_grad_norm = 0.5

94
motrix_rl/src/motrix_rl/tasks/dm_manipulator.py Normal file
View File

@@ -0,0 +1,94 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-manipulator-bring-ball", backend="jax")
@dataclass
class ManipulatorSkrlPpoJax(SkrlCfg):
"""Manipulator SKRL configuration with nested structure (JAX)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 4
agent.learning_rate = 3e-4
agent.ratio_clip = 0.2
agent.entropy_loss_scale = 1e-3
agent.grad_norm_clip = 1.0
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-manipulator-bring-ball", backend="torch")
@dataclass
class ManipulatorSkrlPpoTorch(SkrlCfg):
"""Manipulator SKRL configuration with nested structure (PyTorch)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures (PyTorch - different network)
models.policy.hiddens = [256, 256]
models.value.hiddens = [256, 256]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 4
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 20000
@rlcfg("dm-manipulator-bring-ball")
@dataclass
class ManipulatorRslrlPpo(RslrlCfg):
"""Manipulator RSLRL configuration"""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 500
runner.num_steps_per_env = 24
runner.experiment_name = "dm_manipulator"
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
algo.learning_rate = 3e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 4

94
motrix_rl/src/motrix_rl/tasks/dm_quadruped.py Normal file
View File

@@ -0,0 +1,94 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("dm-quadruped-walk", backend="jax")
@rlcfg("dm-quadruped-run", backend="jax")
@rlcfg("dm-quadruped-escape", backend="jax")
@rlcfg("dm-quadruped-fetch", backend="jax")
@dataclass
class QuadrupedSkrlPpoJax(SkrlCfg):
"""DM quadruped tasks - SKRL PPO configuration (JAX)."""
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.clip_actions = False
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
trainer.timesteps = 20000
@rlcfg("dm-quadruped-walk", backend="torch")
@rlcfg("dm-quadruped-run", backend="torch")
@rlcfg("dm-quadruped-escape", backend="torch")
@rlcfg("dm-quadruped-fetch", backend="torch")
@dataclass
class QuadrupedSkrlPpoTorch(SkrlCfg):
"""DM quadruped tasks - SKRL PPO configuration (PyTorch)."""
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.policy.clip_actions = False
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
trainer.timesteps = 27000
class rslrl:
@rlcfg("dm-quadruped-walk")
@rlcfg("dm-quadruped-run")
@rlcfg("dm-quadruped-escape")
@dataclass
class QuadrupedRslrlPpo(RslrlCfg):
"""DM quadruped tasks - RSLRL PPO configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 1667
runner.num_steps_per_env = 24
runner.experiment_name = "dm_quadruped"
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 32

90
motrix_rl/src/motrix_rl/tasks/dm_reacher.py Normal file
View File

@@ -0,0 +1,90 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
@rlcfg("dm-reacher")
@dataclass
class ReacherSkrlPpo(SkrlCfg):
"""Reacher SKRL configuration with nested structure (JAX/PyTorch)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures (JAX backend)
models.policy.hiddens = [32, 32, 32]
models.value.hiddens = [32, 32, 32]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
trainer.timesteps = 3000
@rlcfg("dm-reacher")
@dataclass
class ReacherSkrlPpoTorch(SkrlCfg):
"""Reacher SKRL configuration with nested structure (PyTorch)."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures (PyTorch backend - larger network)
models.policy.hiddens = [64, 64, 64]
models.value.hiddens = [64, 64, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
# Configure training parameters
trainer.timesteps = 3000
@rlcfg("dm-reacher")
@dataclass
class ReacherRslrlPpo(RslrlCfg):
"""Reacher RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 100
runner.num_steps_per_env = 24
runner.experiment_name = "dm_reacher"
runner.actor.hidden_dims = [32, 32, 32]
runner.critic.hidden_dims = [32, 32, 32]
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 4

171
motrix_rl/src/motrix_rl/tasks/dm_walker.py Normal file
View File

@@ -0,0 +1,171 @@
# 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.
# ==============================================================================
"""Configuration for dm-walker, dm-stander, and dm-runner tasks.
These three tasks share similar configurations:
- JAX: dm-walker, dm-stander, dm-runner all use the same config
- Torch: dm-walker and dm-stander share one config, dm-runner has a different config
- RSLRL: Each task has its own config
"""
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
# =============================================================================
# SKRL JAX Configurations
# =============================================================================
@rlcfg("dm-walker", backend="jax")
@rlcfg("dm-stander", backend="jax")
@rlcfg("dm-runner", backend="jax")
@dataclass
class DmRunnerSkrlJaxCfg(SkrlCfg):
"""Shared SKRL JAX configuration for dm-walker, dm-stander, and dm-runner."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
runner.seed = 42
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 4
agent.learning_rate = 2e-4
# Configure training parameters
# trainer.timesteps = max_env_steps / num_envs
runner.trainer.timesteps = 20000
# =============================================================================
# SKRL Torch Configurations
# =============================================================================
@rlcfg("dm-walker", backend="torch")
@rlcfg("dm-stander", backend="torch")
@dataclass
class DmWalkerStanderSkrlTorchCfg(SkrlCfg):
"""Shared SKRL Torch configuration for dm-walker and dm-stander."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
runner.seed = 42
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 2e-4
# Configure training parameters
runner.trainer.timesteps = 20000
@rlcfg("dm-runner", backend="torch")
@dataclass
class DmRunnerSkrlTorchCfg(SkrlCfg):
"""dm-runner SKRL Torch configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
runner.seed = 42
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 24
agent.learning_epochs = 2
agent.mini_batches = 32
agent.learning_rate = 2e-4
runner.trainer.timesteps = 20000
# =============================================================================
# RSLRL Configurations
# =============================================================================
@rlcfg("dm-walker")
@dataclass
class WalkerRslrlPpo(RslrlCfg):
"""dm-walker RSLRL configuration.
Note: max_iterations = max_env_steps / num_envs / roll_out
= 1024 * 40000 / 2048 / 24 ≈ 833
"""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 833
runner.num_steps_per_env = 24
runner.experiment_name = "dm_walker"
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 4
@rlcfg("dm-stander")
@dataclass
class StanderRslrlPpo(RslrlCfg):
"""dm-stander RSLRL configuration.
Note: max_iterations = max_env_steps / num_envs / roll_out
= 1024 * 40000 / 2048 / 24 ≈ 833
"""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 833
runner.num_steps_per_env = 24
runner.experiment_name = "dm_stander"
algo.learning_rate = 2e-4
algo.num_learning_epochs = 4
algo.num_mini_batches = 32
@rlcfg("dm-runner")
@dataclass
class RunnerRslrlPpo(RslrlCfg):
"""dm-runner RSLRL configuration.
Note: max_iterations = max_env_steps / num_envs / roll_out
= 1024 * 40000 / 2048 / 24 ≈ 833
"""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 833
runner.num_steps_per_env = 24
runner.experiment_name = "dm_runner"
algo.learning_rate = 2e-4
algo.num_learning_epochs = 2
algo.num_mini_batches = 32

103
motrix_rl/src/motrix_rl/tasks/franka_lift_cube.py Normal file
View File

@@ -0,0 +1,103 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("franka-lift-cube", "jax")
@dataclass
class FrankaLiftPPOJax(SkrlCfg):
"""Franka lift cube - SKRL JAX PPO configuration."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 4
agent.mini_batches = 32
agent.learning_rate = 1e-3
# Configure training parameters
trainer.timesteps = 100000
@rlcfg("franka-lift-cube", "torch")
@dataclass
class FrankaLiftPPOTorch(SkrlCfg):
"""Franka lift cube - SKRL Torch PPO configuration."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 8
agent.mini_batches = 4
agent.learning_rate = 3e-4
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.01}
agent.entropy_loss_scale = 0.001
agent.rewards_shaper_scale = 0.01
# Configure training parameters
trainer.timesteps = 100000
class rslrl:
@rlcfg("franka-lift-cube")
@dataclass
class FrankaLiftRslrlPpo(RslrlCfg):
"""Franka lift cube - RSLRL PPO configuration."""
def __post_init__(self):
"""Configure RSLRL runner and algorithm settings."""
runner = self.runner
algo = runner.algorithm
# Runner settings
runner.seed = 42
runner.max_iterations = 500
runner.num_steps_per_env = 64
runner.experiment_name = "franka_lift_cube"
# Network architecture
runner.actor.hidden_dims = [256, 128, 128]
runner.critic.hidden_dims = [256, 128, 128]
# Algorithm parameters (match SKRL Torch config)
algo.learning_rate = 5e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 4
algo.entropy_coef = 1e-3

77
motrix_rl/src/motrix_rl/tasks/franka_open_cabinet.py Normal file
View File

@@ -0,0 +1,77 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("franka-open-cabinet")
@dataclass
class FrankaOpenCabinetPPO(SkrlCfg):
"""Franka open cabinet - SKRL PPO configuration."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 16
agent.learning_epochs = 5
agent.mini_batches = 32
agent.learning_rate = 3e-4
agent.entropy_loss_scale = 0.001
agent.rewards_shaper_scale = 1e-1
# Configure training parameters
runner.seed = 64
trainer.timesteps = 24000
class rslrl:
@rlcfg("franka-open-cabinet")
@dataclass
class FrankaOpenCabinetRslrlPpo(RslrlCfg):
"""Franka open cabinet - RSLRL PPO configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
# Runner settings
runner.seed = 64
runner.max_iterations = 1500
runner.num_steps_per_env = 16
runner.experiment_name = "franka_open_cabinet"
# Network architecture
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
# Algorithm parameters
algo.learning_rate = 3e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 32
algo.entropy_coef = 0.001

125
motrix_rl/src/motrix_rl/tasks/go1.py Normal file
View File

@@ -0,0 +1,125 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("go1-flat-terrain-walk")
@dataclass
class Go1WalkFlatSkrlPpo(SkrlCfg):
"""Go1 robot walk on flat terrain - SKRL PPO configuration."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures (medium size network)
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 5
agent.mini_batches = 3
agent.learning_rate = 3e-4
# Configure training parameters
trainer.timesteps = 30000
@rlcfg("go1-rough-terrain-walk")
@dataclass
class Go1WalkRoughSkrlPpo(Go1WalkFlatSkrlPpo):
"""Go1 robot walk on rough terrain - SKRL PPO configuration.
Uses larger network than flat terrain for more complex terrain handling.
"""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
# Configure model architectures (larger network for rough terrain)
models.policy.hiddens = [512, 256, 128]
models.value.hiddens = [512, 256, 128]
@rlcfg("go1-stairs-terrain-walk")
@dataclass
class Go1WalkStairsPPO(Go1WalkRoughSkrlPpo): ...
class rslrl:
@rlcfg("go1-flat-terrain-walk")
@dataclass
class Go1WalkFlatRslrlPpo(RslrlCfg):
"""Go1 robot walk on flat terrain - RSLRL PPO configuration."""
def __post_init__(self):
"""Configure RSLRL runner and algorithm settings."""
runner = self.runner
algo = runner.algorithm
# Runner settings
runner.seed = 42
runner.max_iterations = 1000
runner.num_steps_per_env = 24
runner.experiment_name = "go1_flat_terrain_walk"
# Network architecture (medium size for flat terrain)
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
# Algorithm parameters (match SKRL config)
algo.learning_rate = 3e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 3
@rlcfg("go1-rough-terrain-walk")
@dataclass
class Go1WalkRoughRslrlPpo(Go1WalkFlatRslrlPpo):
"""Go1 robot walk on rough terrain - RSLRL PPO configuration.
Uses larger network than flat terrain for more complex terrain handling.
"""
def __post_init__(self):
"""Override network architecture for rough terrain."""
super().__post_init__()
# Override experiment name
self.runner.experiment_name = "go1_rough_terrain_walk"
# Override network architecture (larger for rough terrain)
self.runner.actor.hidden_dims = [512, 256, 128]
self.runner.critic.hidden_dims = [512, 256, 128]
@rlcfg("go1-stairs-terrain-walk")
@dataclass
class Go1WalkStairsRslrlPpo(Go1WalkRoughRslrlPpo):
"""Go1 robot walk on stairs terrain - RSLRL PPO configuration.
Uses same configuration as rough terrain since stairs also require complex handling.
"""
def __post_init__(self):
super().__post_init__()
self.runner.experiment_name = "go1_stairs_terrain_walk"

74
motrix_rl/src/motrix_rl/tasks/go2.py Normal file
View File

@@ -0,0 +1,74 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("go2-flat-terrain-walk")
@dataclass
class Go2WalkFlatSkrlPpo(SkrlCfg):
"""Go2 robot walk on flat terrain - SKRL PPO configuration."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# Configure model architectures (medium size network)
models.policy.hiddens = [256, 128, 64]
models.value.hiddens = [256, 128, 64]
# Configure PPO agent parameters
agent.rollouts = 24
agent.learning_epochs = 5
agent.mini_batches = 3
agent.learning_rate = 3e-4
# Configure training parameters
trainer.timesteps = 30000
class rslrl:
@rlcfg("go2-flat-terrain-walk")
@dataclass
class Go2WalkFlatRslrlPpo(RslrlCfg):
"""Go2 robot walk on flat terrain - RSLRL PPO configuration."""
def __post_init__(self):
"""Configure RSLRL runner and algorithm settings."""
runner = self.runner
algo = runner.algorithm
# Runner settings
runner.seed = 42
runner.max_iterations = 1000
runner.num_steps_per_env = 24
runner.experiment_name = "go2_flat_terrain_walk"
# Network architecture (medium size for flat terrain)
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
# Algorithm parameters
algo.learning_rate = 3e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 3

69
motrix_rl/src/motrix_rl/tasks/pendulum.py Normal file
View File

@@ -0,0 +1,69 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
# -- docs-tag-start: pendulum-train-cfg --
@rlcfg("pendulum")
@dataclass
class PendulumSkrlPpo(SkrlCfg):
"""Pendulum SKRL configuration with nested structure."""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
runner = self.runner
# Configure model architectures
runner.models.policy.hiddens = [64, 64]
runner.models.value.hiddens = [64, 64]
# Configure PPO agent parameters
agent = runner.agent
agent.rollouts = 32
agent.learning_epochs = 5
agent.mini_batches = 4
agent.learning_rate = 3e-4
# Configure training parameters
# trainer.timesteps = max_env_steps / num_envs
runner.trainer.timesteps = 5000
# -- docs-tag-end: pendulum-train-cfg --
@rlcfg("pendulum")
@dataclass
class PendulumRslrlPpo(RslrlCfg):
"""Pendulum RSLRL configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 42
runner.max_iterations = 150
runner.num_steps_per_env = 32
runner.experiment_name = "pendulum"
runner.actor.hidden_dims = [64, 64]
runner.critic.hidden_dims = [64, 64]
algo.learning_rate = 3e-4
algo.entropy_coef = 0.005
algo.num_learning_epochs = 5
algo.num_mini_batches = 4

82
motrix_rl/src/motrix_rl/tasks/point_mass.py Normal file
View File

@@ -0,0 +1,82 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.skrl.config import SkrlCfg
def _configure_point_mass_runner(runner) -> None:
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.separate = True
models.policy.hiddens = [32, 32]
models.value.hiddens = [32, 32]
agent.rollouts = 16
agent.learning_epochs = 4
agent.mini_batches = 4
trainer.timesteps = 20000
@rlcfg("point_mass", backend="jax")
@dataclass
class PointMassSkrlPpoJax(SkrlCfg):
"""Point mass SKRL configuration for the JAX backend."""
def __post_init__(self):
self.num_envs = 256
runner = self.runner
agent = runner.agent
_configure_point_mass_runner(runner)
agent.learning_rate = 1e-3
agent.entropy_loss_scale = 0.01
agent.rewards_shaper_scale = 0.05
agent.grad_norm_clip = 0.1
agent.discount_factor = 0.995
agent.lam = 0.97
agent.value_loss_scale = 0.5
agent.value_clip = 10.0
agent.clip_predicted_values = False
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.02}
agent.kl_threshold = 0.03
@rlcfg("point_mass", backend="torch")
@dataclass
class PointMassSkrlPpoTorch(SkrlCfg):
"""Point mass SKRL configuration for the PyTorch backend."""
def __post_init__(self):
runner = self.runner
agent = runner.agent
models = runner.models
models.separate = False
models.policy.hiddens = [32, 32]
models.value.hiddens = [32, 32]
agent.rollouts = 32
agent.learning_epochs = 5
agent.mini_batches = 4
runner.trainer.timesteps = 3500
agent.learning_rate = 3e-4
agent.entropy_loss_scale = 0.1

74
motrix_rl/src/motrix_rl/tasks/rm65_open_cabinet.py Normal file
View File

@@ -0,0 +1,74 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("rm65-open-cabinet", "torch")
@dataclass
class RM65OpenCabinetPPOTorch(SkrlCfg):
"""RM65 open cabinet - SKRL Torch PPO configuration."""
def __post_init__(self):
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
models.separate = True
models.policy.hiddens = [512, 256, 128]
models.value.hiddens = [512, 256, 128]
agent.rollouts = 32
agent.learning_epochs = 5
agent.mini_batches = 32
agent.learning_rate = 8e-5
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.005}
agent.ratio_clip = 0.10
agent.grad_norm_clip = 0.4
agent.entropy_loss_scale = 2e-4
agent.rewards_shaper_scale = 5e-2
runner.seed = 64
trainer.timesteps = 21000
class rslrl:
@rlcfg("rm65-open-cabinet")
@dataclass
class RM65OpenCabinetRslrlPpo(RslrlCfg):
"""RM65 open cabinet - RSLRL PPO configuration."""
def __post_init__(self):
runner = self.runner
algo = runner.algorithm
runner.seed = 64
runner.max_iterations = 1500
runner.num_steps_per_env = 24
runner.experiment_name = "rm65_open_cabinet"
runner.actor.hidden_dims = [256, 128, 64]
runner.critic.hidden_dims = [256, 128, 64]
algo.learning_rate = 3e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 8
algo.entropy_coef = 0.001

110
motrix_rl/src/motrix_rl/tasks/shadow_hand_repose.py Normal file
View File

@@ -0,0 +1,110 @@
# 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.
# ==============================================================================
from dataclasses import dataclass
from motrix_rl.registry import rlcfg
from motrix_rl.rslrl.cfg import RslrlCfg
from motrix_rl.skrl.config import SkrlCfg
class skrl:
@rlcfg("shadow-hand-repose")
@dataclass
class ShadowHandReposePPO(SkrlCfg):
"""Shadow Hand Repose PPO configuration.
Configuration for training Shadow Hand to reach target hand pose.
Uses large-scale parallel training with 8192 environments.
"""
def __post_init__(self):
"""Configure nested SKRL runner settings."""
self.num_envs = 8192
runner = self.runner
models = runner.models
agent = runner.agent
trainer = runner.trainer
# ===== Basic Settings =====
runner.seed = 42
# ===== Network Architecture =====
models.policy.hiddens = [512, 512, 256, 128]
models.value.hiddens = [512, 512, 256, 128]
# ===== PPO Core Parameters =====
agent.rollouts = 16
agent.learning_epochs = 5
agent.mini_batches = 4
agent.discount_factor = 0.99
agent.lam = 0.95
# ===== Learning Rate =====
agent.learning_rate = 5.0e-04
agent.learning_rate_scheduler_kwargs = {"kl_threshold": 0.016}
# ===== Clipping =====
agent.ratio_clip = 0.2
agent.value_clip = 0.2
agent.clip_predicted_values = True
agent.grad_norm_clip = 1.0
# ===== Loss Coefficients =====
agent.entropy_loss_scale = 0.0
agent.value_loss_scale = 2.0
agent.kl_threshold = 0.0
# ===== Reward Shaping =====
agent.rewards_shaper_scale = 0.01
# ===== Training Control =====
agent.random_timesteps = 0
agent.learning_starts = 0
agent.time_limit_bootstrap = False
# ===== Training Parameters =====
trainer.timesteps = 24000
class rslrl:
@rlcfg("shadow-hand-repose")
@dataclass
class ShadowHandReposeRslrlPpo(RslrlCfg):
"""Shadow Hand Repose RSLRL PPO configuration."""
def __post_init__(self):
"""Configure RSLRL runner and algorithm settings."""
self.num_envs = 8192
runner = self.runner
algo = runner.algorithm
# ===== Basic Settings =====
runner.seed = 42
# max_iterations = max_env_steps / num_envs / roll_out = 200000000 / 8192 / 16 ≈ 1525
runner.max_iterations = 1500
runner.num_steps_per_env = 16
runner.experiment_name = "shadow_hand_repose"
# ===== Network Architecture =====
runner.actor.hidden_dims = [512, 512, 256, 128]
runner.critic.hidden_dims = [512, 512, 256, 128]
# ===== Algorithm Parameters =====
algo.learning_rate = 5.0e-4
algo.num_learning_epochs = 5
algo.num_mini_batches = 4
algo.entropy_coef = 0.0

149
motrix_rl/src/motrix_rl/utils.py
View File

@@ -13,7 +13,11 @@
# limitations under the License.
# ==============================================================================
import dataclasses
from dataclasses import dataclass
from typing import Any, TypeVar
T = TypeVar("T")
@dataclass
@@ -59,3 +63,148 @@ def get_device_supports() -> DeviceSupports:
pass
return supports
def class_to_dict(obj) -> dict | list | Any:
"""Recursively convert a dataclass to a dictionary.
Args:
obj: The object to convert (dataclass, list, dict, or primitive)
Returns:
Dictionary representation with nested dataclasses recursively converted
"""
if dataclasses.is_dataclass(obj):
return {k: class_to_dict(v) for k, v in dataclasses.asdict(obj).items()}
elif isinstance(obj, list):
return [class_to_dict(item) for item in obj]
elif isinstance(obj, dict):
return {k: class_to_dict(v) for k, v in obj.items()}
else:
return obj
def cfg_override(cfg: T, overrides: dict[str, Any]) -> T:
"""Override dataclass fields using dot-notation path keys.
This function creates a new dataclass instance with specified field values
overridden, leaving the original config unchanged. Nested dataclasses are
handled using dot notation in the key path.
Args:
cfg: The original dataclass configuration object
overrides: Dictionary with path keys (e.g., "runner.seed", "num_envs")
where each key is a dot-separated path to the field to override
Returns:
A new dataclass instance with overrides applied
Raises:
KeyError: If a path key is invalid or references a non-existent field
TypeError: If an intermediate field is not a dataclass or if a value
type doesn't match the expected field type
Examples:
>>> from motrix_rl.rslrl.cfg import RslrlCfg
>>> base_cfg = RslrlCfg()
>>> overrides = {
... "num_envs": 4096,
... "runner.seed": 123,
... "runner.algorithm.num_learning_epochs": 10,
... }
>>> new_cfg = cfg_override(base_cfg, overrides)
>>> assert new_cfg.num_envs == 4096
>>> assert new_cfg.runner.seed == 123
>>> assert new_cfg.runner.algorithm.num_learning_epochs == 10
"""
if not overrides:
return cfg
if not dataclasses.is_dataclass(cfg):
raise TypeError(f"cfg must be a dataclass, got {type(cfg).__name__}")
# Group overrides by their parent paths to apply them efficiently
# Structure: {parent_path: {field_name: value}}
# For "num_envs": parent_path=[], field_name="num_envs"
# For "runner.seed": parent_path=["runner"], field_name="seed"
override_tree: dict[tuple[str, ...], dict[str, Any]] = {}
for key, value in overrides.items():
parts = key.split(".")
if len(parts) == 1:
# Top-level field
parent_path = tuple()
field_name = parts[0]
else:
# Nested field
parent_path = tuple(parts[:-1])
field_name = parts[-1]
if parent_path not in override_tree:
override_tree[parent_path] = {}
override_tree[parent_path][field_name] = value
# Apply overrides from deepest to shallowest to minimize object copies
sorted_paths = sorted(override_tree.keys(), key=lambda p: len(p), reverse=True)
def apply_overrides_at_path(
obj: Any, path: tuple[str, ...], field_overrides: dict[str, Any], parent_path: tuple[str, ...] = ()
) -> Any:
"""Apply overrides to an object at a specific path.
Args:
obj: The current object (dataclass or primitive)
path: Tuple of field names to navigate through
field_overrides: Dict of field names to values to apply at the target
parent_path: The full path from the root (for error messages)
Returns:
New object with overrides applied
"""
if not path:
# We're at the target - apply the overrides
if not dataclasses.is_dataclass(obj):
# Build the full path for the error message
full_path_parts = list(parent_path) if parent_path else ["(root)"]
raise TypeError(
f"Cannot navigate into non-dataclass field '{full_path_parts[-1]}' of type {type(obj).__name__}"
)
# Validate field names exist - use set difference for efficiency
obj_fields = {f.name for f in dataclasses.fields(obj)}
invalid_fields = set(field_overrides.keys()) - obj_fields
if invalid_fields:
raise KeyError(
f"Invalid fields {sorted(invalid_fields)} for {type(obj).__name__}. "
f"Valid fields: {sorted(obj_fields)}"
)
# Use dataclasses.replace to create a new instance with overrides
return dataclasses.replace(obj, **field_overrides)
# Need to navigate deeper - recurse to modify nested dataclass
if not dataclasses.is_dataclass(obj):
raise TypeError(f"Cannot navigate into non-dataclass field '{path[0]}' of type {type(obj).__name__}")
# Validate the navigation path exists
obj_fields = {f.name: f for f in dataclasses.fields(obj)}
if path[0] not in obj_fields:
valid_fields = sorted(obj_fields.keys())
raise KeyError(f"Invalid path component '{path[0]}' for {type(obj).__name__}. Valid fields: {valid_fields}")
# Get the nested object
nested_obj = getattr(obj, path[0])
# Recursively apply overrides to the nested object
new_nested_obj = apply_overrides_at_path(nested_obj, path[1:], field_overrides, parent_path + (path[0],))
# Return a new instance of the current object with the nested field replaced
return dataclasses.replace(obj, **{path[0]: new_nested_obj})
# Start with the original cfg and apply each group of overrides
result = cfg
for path in sorted_paths:
field_overrides = override_tree[path]
result = apply_overrides_at_path(result, path, field_overrides, parent_path=path)
return result

85
motrix_rl/template/rslrl_config.yaml Normal file
View File

@@ -0,0 +1,85 @@
runner:
class_name: OnPolicyRunner
# General
num_steps_per_env: 24 # Number of steps per environment per iteration
max_iterations: 1500 # Number of policy updates
seed: 1
# Observations
obs_groups: { "actor": ["policy"], "critic": ["policy", "privileged"] } # Maps from observation sets to groups. See `vec_env.py` for more information
# Logging parameters
save_interval: 50 # Check for potential saves every `save_interval` iterations
experiment_name: walking_experiment
run_name: ""
# Logging writer
logger: tensorboard # tensorboard, neptune, wandb
neptune_project: rsl_rl
wandb_project: rsl_rl
# Actor
actor:
class_name: MLPModel
hidden_dims: [256, 256, 256]
activation: elu
obs_normalization: false
stochastic: true
init_noise_std: 1.0
noise_std_type: "scalar" # 'scalar' or 'log'
state_dependent_std: false
# Critic
critic:
class_name: MLPModel
hidden_dims: [256, 256, 256]
activation: elu
obs_normalization: false
stochastic: false
# Algorithm
algorithm:
class_name: PPO
# Training
optimizer: adam # adam, adamw, sgd, rmsprop
learning_rate: 0.001
num_learning_epochs: 5
num_mini_batches: 4 # mini batch size = num_envs * num_steps / num_mini_batches
schedule: adaptive # adaptive, fixed
# Value function
value_loss_coef: 1.0
clip_param: 0.2
use_clipped_value_loss: true
# Surrogate loss
desired_kl: 0.01
entropy_coef: 0.01
gamma: 0.99
lam: 0.95
max_grad_norm: 1.0
# Miscellaneous
normalize_advantage_per_mini_batch: false
# Random network distillation
rnd_cfg:
weight: 0.0 # Initial weight of the RND reward
weight_schedule: null # This is a dictionary with a required key called "mode". Please check the RND module for more information
reward_normalization: false # Whether to normalize RND reward
# Learning parameters
learning_rate: 0.001 # Learning rate for RND
# Network parameters
num_outputs: 1 # Number of outputs of RND network. Note: if -1, then the network will use dimensions of the observation
predictor_hidden_dims: [-1] # Hidden dimensions of predictor network
target_hidden_dims: [-1] # Hidden dimensions of target network
# Symmetry augmentation
symmetry_cfg:
use_data_augmentation: true # This adds symmetric trajectories to the batch
use_mirror_loss: false # This adds symmetry loss term to the loss function
data_augmentation_func: null # String containing the module and function name to import
# Example: "legged_gym.envs.locomotion.anymal_c.symmetry:get_symmetric_states"
#
# .. code-block:: python
#
# @torch.no_grad()
# def get_symmetric_states(
# env: VecEnv, obs: Optional[torch.Tensor] = None, actions: Optional[torch.Tensor] = None,
# ) -> Tuple[torch.Tensor, torch.Tensor]:
#
mirror_loss_coeff: 0.0 # Coefficient for symmetry loss term. If 0, no symmetry loss is used

75
motrix_rl/template/skrl_config.yaml Normal file
View File

@@ -0,0 +1,75 @@
seed: 42
# Models are instantiated using skrl's model instantiator utility
# https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html
models:
separate: False
policy: # gaussian model
class: "GaussianMixin"
clip_actions: True
clip_log_std: True
initial_log_std: 0.0
min_log_std: -20.0
max_log_std: 2.0
input: "STATES"
hiddens: [32, 32]
hidden_activation: ["elu", "elu"]
output: "ACTIONS"
output_activation: "tanh"
output_scale: 1.0
value: # deterministic model
class: "DeterministicMixin"
clip_actions: False
input: "STATES"
hiddens: [32, 32]
hidden_activation: ["elu", "elu"]
output: "ONE"
output_activation: ""
output_scale: 1.0
# Memory
# https://skrl.readthedocs.io/en/latest/api/memories/random.html
memory:
class: "RandomMemory"
memory_size: -1 # -1: automatically determined value
# PPO agent configuration (field names are from PPO_DEFAULT_CONFIG)
# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html
agent:
class: "PPO"
rollouts: 16
learning_epochs: 8
mini_batches: 1
discount_factor: 0.99
lambda: 0.95
learning_rate: 3.e-4
learning_rate_scheduler: "KLAdaptiveLR"
learning_rate_scheduler_kwargs:
kl_threshold: 0.008
random_timesteps: 0 # random exploration steps
learning_starts: 0 # learning starts after this many steps
grad_norm_clip: 1.0
ratio_clip: 0.2
value_clip: 0.2
clip_predicted_values: True
entropy_loss_scale: 0.0
value_loss_scale: 2.0
kl_threshold: 0
rewards_shaper_scale: 1.0
time_limit_bootstrap: False
# logging and checkpoint
experiment:
directory: "runs"
experiment_name: ""
write_interval: 16
checkpoint_interval: 80
store_separately: False
wandb: False
wandb_kwargs: null
# Sequential trainer
# https://skrl.readthedocs.io/en/latest/api/trainers/sequential.html
trainer:
class: "SequentialTrainer"
timesteps: 1600
environment_info: "log"

158
motrix_rl/tests/test_utils.py Normal file
View File

@@ -0,0 +1,158 @@
# 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.
# ==============================================================================
"""Tests for motrix_rl.utils.cfg_override function."""
import dataclasses
from dataclasses import dataclass
import pytest
from motrix_rl.utils import cfg_override
@dataclass
class NestedConfig:
value: int = 10
name: str = "default"
@dataclass
class MiddleConfig:
nested: NestedConfig = dataclasses.field(default_factory=NestedConfig)
flag: bool = True
@dataclass
class RootConfig:
middle: MiddleConfig = dataclasses.field(default_factory=MiddleConfig)
count: int = 5
label: str = "root"
class TestCfgOverride:
"""Tests for cfg_override function."""
def test_nested_overrides(self):
"""Test overriding fields at all nesting levels."""
cfg = RootConfig()
# Top-level
result = cfg_override(cfg, {"count": 42})
assert result.count == 42
assert result.label == "root"
# One-level nested
result = cfg_override(cfg, {"middle.flag": False})
assert result.middle.flag is False
# Deep nested
result = cfg_override(cfg, {"middle.nested.value": 99})
assert result.middle.nested.value == 99
def test_multiple_overrides(self):
"""Test overriding multiple fields at different levels."""
cfg = RootConfig()
overrides = {
"count": 42,
"label": "modified",
"middle.flag": False,
"middle.nested.name": "custom",
}
result = cfg_override(cfg, overrides)
assert result.count == 42
assert result.label == "modified"
assert result.middle.flag is False
assert result.middle.nested.name == "custom"
def test_type_and_list_fields(self):
"""Test overriding different field types."""
cfg = RootConfig()
# String
result = cfg_override(cfg, {"label": "new_label"})
assert result.label == "new_label"
# Bool
result = cfg_override(cfg, {"middle.flag": False})
assert result.middle.flag is False
# Float (Python dataclasses don't enforce types)
result = cfg_override(cfg, {"count": 3.14})
assert result.count == 3.14
# List
result = cfg_override(cfg, {"count": [1, 2, 3]})
assert result.count == [1, 2, 3]
def test_immutability_and_empty_overrides(self):
"""Test immutability and empty overrides."""
cfg = RootConfig()
original_count = cfg.count
# Empty overrides returns original
result = cfg_override(cfg, {})
assert result is cfg
# Original unchanged
result = cfg_override(cfg, {"count": 42})
assert cfg.count == original_count
assert result.count == 42
def test_error_handling(self):
"""Test error handling for invalid inputs."""
cfg = RootConfig()
# Non-existent top-level field
with pytest.raises(KeyError, match="Invalid field"):
cfg_override(cfg, {"nonexistent": 1})
# Non-existent nested field
with pytest.raises(KeyError, match="Invalid field"):
cfg_override(cfg, {"middle.nonexistent": 1})
# Non-existent intermediate path
with pytest.raises(KeyError, match="Invalid path component"):
cfg_override(cfg, {"invalid.path": 1})
# Navigate into non-dataclass field
with pytest.raises(TypeError, match="Cannot navigate into non-dataclass"):
cfg_override(cfg, {"count.something": 1})
# Non-dataclass cfg
with pytest.raises(TypeError, match="cfg must be a dataclass"):
cfg_override({"foo": "bar"}, {"foo": "baz"})
def test_rslrl_config(self):
"""Test with actual RSLRL configuration classes."""
from motrix_rl.rslrl.cfg import RslrlCfg
cfg = RslrlCfg()
overrides = {
"num_envs": 4096,
"runner.seed": 123,
"runner.algorithm.num_learning_epochs": 10,
"runner.algorithm.learning_rate": 1e-4,
"runner.actor.hidden_dims": [512, 256, 128],
}
result = cfg_override(cfg, overrides)
assert result.num_envs == 4096
assert result.play_num_envs == 16 # Unchanged
assert result.runner.seed == 123
assert result.runner.algorithm.num_learning_epochs == 10
assert result.runner.algorithm.learning_rate == 1e-4
assert result.runner.actor.hidden_dims == [512, 256, 128]