chore: release v0.0.1

motrix_rl/README.md

@@ -0,0 +1,3 @@
+# MotrixLab
+
+集成 rl 框架，并针对 motrix_envs 里的各种环境提供默认的训练参数。 预期会支持不同的 rl 框架的不同 rl 算法，但目前只支持 了 skrl 框架的 ppo 算法

motrix_rl/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["uv_build>=0.9.6,<0.10.0"]
+build-backend = "uv_build"
+
+[project]
+name = "motrix-rl"
+version = "0.0.1"
+description = "Reinforcement learning training framework based on SKRL with multi-backend unified training interface."
+authors = [{ name = "Motphys", email = "developers@motphys.com" }]
+requires-python = "==3.10.*"
+readme = "README.md"
+license = "Apache-2.0"
+dependencies = ["gymnasium===1.1.1", "python-abc>=0.2.0", "motrix-envs"]
+
+[project.optional-dependencies]
+skrl-jax = [
+    "skrl===1.4.3; sys_platform == 'linux'",
+    "jax[cuda12]==0.4.34; sys_platform == 'linux'",
+    "flax===0.10.4; sys_platform == 'linux'",
+]
+skrl-torch = [
+    "skrl===1.4.3",
+    "torch===2.7.0+cu128",
+    "torchvision===v0.22.0+cu128",
+    "torchaudio===2.7.0+cu128",
+]
+test = ["pytest~=8.3.3"]
+
+[tool.uv.sources]
+motrix-envs = { workspace = true }

motrix_rl/src/motrix_rl/__init__.py

@@ -0,0 +1,16 @@
+# 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 cfgs, skrl  # noqa: F401

motrix_rl/src/motrix_rl/base.py

@@ -0,0 +1,42 @@
+# 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

motrix_rl/src/motrix_rl/cfgs.py

@@ -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.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("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
+
+
+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 * 60000
+        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

motrix_rl/src/motrix_rl/registry.py

@@ -0,0 +1,114 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import logging
+from dataclasses import dataclass, field
+from typing import 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)
+
+
+@dataclass
+class EnvRlCfgs:
+    cfgs: dict[str, dict[str, Type[BaseRLCfg]]] = field(default_factory=dict)
+    """
+    The RL configuration classes available for this environment.
+    Structure: {rl_framework: {backend: config_class}}
+    Example: {"skrl": {"jax": JaxConfig, "torch": TorchConfig}}
+    """
+
+
+# RL configuration registry. map from env name to EnvMeta
+_rlcfgs: dict[str, EnvRlCfgs] = {}
+
+
+def _register_rlcfg(env_name: str, rllib: str, backend: str, train_cfg_cls: Type[BaseRLCfg]):
+    """
+    Register a training configuration class for an environment, reinforcement learning framework, and backend.
+
+    Args:
+        env_name: Environment name
+        rllib: RL framework name (e.g., "skrl")
+        backend: Backend name (e.g., "jax", "torch")
+        train_cfg_cls: Configuration class
+    """
+    if not env_registry.contains(env_name):
+        raise ValueError(f"Environment '{env_name}' is not registered in env_registry.")
+
+    logger.info(f"Registering RL config for env '{env_name}', RL framework '{rllib}', and backend '{backend}'")
+    if env_name not in _rlcfgs:
+        _rlcfgs[env_name] = EnvRlCfgs()
+    if rllib not in _rlcfgs[env_name].cfgs:
+        _rlcfgs[env_name].cfgs[rllib] = {}
+    _rlcfgs[env_name].cfgs[rllib][backend] = train_cfg_cls
+
+
+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.
+
+    Args:
+        env_name: Environment name
+        backend: Backend name (e.g., "jax", "torch"). If None, registers for all backends.
+    """
+
+    def decorator(cls: Type[TRLCfg]) -> Type[TRLCfg]:
+        backends = ["jax", "torch"] if backend is None else [backend]
+        for b in backends:
+            _register_rlcfg(env_name, "skrl", b, cls)
+        return cls
+
+    return decorator
+
+
+def default_rl_cfg(env_name: str, rllib: str, backend: str) -> BaseRLCfg:
+    """
+    Get the default training configuration for an environment, reinforcement learning framework, and backend.
+
+    Args:
+        env_name: Environment name
+        rllib: RL framework name (e.g., "skrl")
+        backend: Backend name (e.g., "jax", "torch")
+
+    Returns:
+        The configuration class instance. Will use backend-specific config if available,
+        otherwise falls back to universal config (backend=None).
+    """
+    if env_name not in _rlcfgs:
+        raise ValueError(f"Environment '{env_name}' is not registered.")
+    meta: EnvRlCfgs = _rlcfgs.get(env_name)
+    if rllib not in meta.cfgs:
+        raise ValueError(f"RL framework '{rllib}' is not supported for environment '{env_name}'.")
+
+    framework_configs = meta.cfgs[rllib]
+
+    # Try to get backend-specific config first
+    if backend in framework_configs:
+        return framework_configs[backend]()
+
+    # Fall back to universal config (backend=None) if backend-specific one is not found
+    if None in framework_configs:
+        return framework_configs[None]()
+
+    # If no universal config exists, raise an error
+    raise ValueError(
+        f"No configuration found for environment '{env_name}', RL framework '{rllib}', backend '{backend}', \
+            and no universal configuration available."
+    )

motrix_rl/src/motrix_rl/skrl/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+# ==============================================================================
+
+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}"

motrix_rl/src/motrix_rl/skrl/cfg.py

@@ -0,0 +1,73 @@
+# 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

motrix_rl/src/motrix_rl/skrl/jax/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+
+
+def _inherits_from(cls, base_class_name):
+    """Check if cls inherits from a class with the given base_class_name."""
+    return any(base.__name__ == base_class_name for base in cls.__mro__)
+
+
+def wrap_env(env, enable_render: bool = False):
+    """Wrap the environment based on its type."""
+    if _inherits_from(env.__class__, "NpEnv"):
+        from motrix_rl.skrl.jax.wrap_np import SkrlNpWrapper
+
+        return SkrlNpWrapper(env, enable_render=enable_render)
+    else:
+        raise ValueError(f"Unsupported environment type: {env.__class__.__name__}")

motrix_rl/src/motrix_rl/skrl/jax/train/ppo.py

@@ -0,0 +1,300 @@
+# 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 typing import Any
+
+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
+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.skrl import get_log_dir
+from motrix_rl.skrl.cfg import PPOCfg
+from motrix_rl.skrl.jax import wrap_env
+
+
+def _get_cfg(
+    rlcfg: PPOCfg,
+    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
+    else:
+        cfg["rewards_shaper"] = None
+    cfg["time_limit_bootstrap"] = rlcfg.time_limit_bootstrap
+    cfg["state_preprocessor"] = RunningStandardScaler
+    cfg["state_preprocessor_kwargs"] = {
+        "size": env.observation_space,
+        "device": env.device,
+    }
+    cfg["value_preprocessor"] = RunningStandardScaler
+    cfg["value_preprocessor_kwargs"] = {"size": 1, "device": env.device}
+    # logging to TensorBoard and write checkpoints (in timesteps)
+    if log_dir:
+        cfg["experiment"]["write_interval"] = rlcfg.check_point_interval
+        cfg["experiment"]["checkpoint_interval"] = rlcfg.check_point_interval
+        cfg["experiment"]["directory"] = log_dir
+    else:
+        cfg["experiment"]["write_interval"] = 0
+        cfg["experiment"]["checkpoint_interval"] = 0
+
+    return cfg
+
+
+class PPO(BasePPO):
+    _total_custom_rewards: dict[str, np.ndarray] = {}
+
+    def record_transition(
+        self,
+        states,
+        actions,
+        rewards,
+        next_states,
+        terminated,
+        truncated,
+        infos,
+        timestep,
+        timesteps,
+    ) -> None:
+        super().record_transition(
+            states,
+            actions,
+            rewards,
+            next_states,
+            terminated,
+            truncated,
+            infos,
+            timestep,
+            timesteps,
+        )
+
+        if "Reward" in infos:
+            for key, value in infos["Reward"].items():
+                self.tracking_data[f"Reward Instant / {key} (max)"].append(jnp.max(value))
+                self.tracking_data[f"Reward Instant / {key} (min)"].append(jnp.min(value))
+                self.tracking_data[f"Reward Instant / {key} (mean)"].append(jnp.mean(value))
+                if key not in self._total_custom_rewards:
+                    self._total_custom_rewards[key] = jnp.zeros_like(value)
+                self._total_custom_rewards[key] += value
+            done = terminated | truncated
+            done = done.reshape(-1)
+            if done.any():
+                for key in self._total_custom_rewards:
+                    self.tracking_data[f"Reward Total/ {key} (mean)"].append(
+                        jnp.mean(self._total_custom_rewards[key][done])
+                    )
+                    self.tracking_data[f"Reward Total/ {key} (min)"].append(
+                        jnp.min(self._total_custom_rewards[key][done])
+                    )
+                    self.tracking_data[f"Reward Total/ {key} (max)"].append(
+                        jnp.max(self._total_custom_rewards[key][done])
+                    )
+
+                    self._total_custom_rewards[key] = self._total_custom_rewards[key] * (1 - done)
+
+        if "metrics" in infos:
+            for key, value in infos["metrics"].items():
+                self.tracking_data[f"metrics / {key} (max)"].append(jnp.max(value))
+                self.tracking_data[f"metrics / {key} (min)"].append(jnp.min(value))
+                self.tracking_data[f"metrics / {key} (mean)"].append(jnp.mean(value))
+
+
+class Trainer:
+    _trainer: SequentialTrainer
+    _env_name: str
+    _sim_backend: str
+    _rlcfg: PPOCfg
+    _enable_render: bool
+
+    def __init__(
+        self,
+        env_name: str,
+        sim_backend: str = None,
+        enable_render: bool = False,
+        cfg_override: dict = None,
+    ) -> None:
+        rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
+        if cfg_override is not None:
+            rlcfg = rlcfg.replace(**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.
+        """
+        rlcfg = self._rlcfg
+        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
+
+        set_seed(rlcfg.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)
+        cfg_trainer = {
+            "timesteps": rlcfg.max_batch_env_steps,
+            "headless": not self._enable_render,
+        }
+        trainer = SequentialTrainer(cfg=cfg_trainer, env=skrl_env, agents=agent)
+        trainer.train()
+
+    def play(self, policy: str) -> None:
+        import time
+
+        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)
+        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)
+        agent.load(policy)
+        obs, _ = env.reset()
+
+        fps = 60
+        while True:
+            t = time.time()
+            outputs = agent.act(obs, timestep=0, timesteps=0)
+            actions = outputs[-1].get("mean_actions", outputs[0])
+            obs, _, _, _, _ = env.step(actions)
+            env.render()
+            delta_time = time.time() - t
+            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
+        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,
+            ):
+                Model.__init__(self, observation_space, action_space, device, **kwargs)
+                GaussianMixin.__init__(
+                    self,
+                    clip_actions,
+                    clip_log_std,
+                    min_log_std,
+                    max_log_std,
+                    reduction,
+                )
+
+            @nn.compact  # marks the given module method allowing inlined submodules
+            def __call__(self, inputs, role):
+                x = inputs["states"]
+                for size in rlcfg.policy_hidden_layer_sizes:
+                    x = nn.elu(nn.Dense(size)(x))
+                x = nn.Dense(self.num_actions)(x)
+                log_std = self.param("log_std", lambda _: jnp.ones(self.num_actions))
+                return x, log_std, {}
+
+        class Value(DeterministicMixin, Model):
+            def __init__(
+                self,
+                observation_space,
+                action_space,
+                device=None,
+                clip_actions=False,
+                **kwargs,
+            ):
+                Model.__init__(self, observation_space, action_space, device, **kwargs)
+                DeterministicMixin.__init__(self, clip_actions)
+
+            @nn.compact  # marks the given module method allowing inlined submodules
+            def __call__(self, inputs, role):
+                x = inputs["states"]
+                for size in rlcfg.value_hidden_layer_sizes:
+                    x = nn.elu(nn.Dense(size)(x))
+                x = nn.Dense(1)(x)
+                return x, {}
+
+        models = {}
+        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,
+        )
+
+        # 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)
+
+        agent = PPO(
+            models=models,
+            memory=memory,
+            cfg=ppo_cfg,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            device=env.device,
+        )
+        return agent
+
+
+class Player:
+    def __init__(self, env_name: str, sim_backend: str = None) -> None:
+        pass

motrix_rl/src/motrix_rl/skrl/jax/wrap_np.py

@@ -0,0 +1,80 @@
+# 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 typing import Any, Tuple
+
+import gymnasium
+import jax
+import numpy as np
+from skrl.envs.jax import Wrapper as SkrlWrapper
+
+from motrix_envs.np.env import NpEnv
+from motrix_envs.np.renderer import NpRenderer
+
+
+class SkrlNpWrapper(SkrlWrapper):
+    """
+    Wrap the numpy-based environment to be compatible with skrl
+    """
+
+    _env: NpEnv
+    _renderer: NpRenderer = None
+
+    def __init__(self, env: NpEnv, enable_render: bool = False):
+        super().__init__(env)
+        if enable_render:
+            self._renderer = NpRenderer(env)
+
+    def reset(self) -> Tuple[jax.Array, Any]:
+        state = self._env.init_state()
+        return state.obs, state.info
+
+    def step(
+        self, actions: jax.Array
+    ) -> Tuple[
+        jax.Array,
+        jax.Array,
+        jax.Array,
+        jax.Array,
+        Any,
+    ]:
+        actions = np.array(actions)
+        state = self._env.step(actions)
+        return (
+            state.obs,
+            state.reward.reshape(-1, 1),
+            state.terminated.reshape(-1, 1),
+            state.truncated.reshape(-1, 1),
+            state.info,
+        )
+
+    def render(self, *args, **kwargs) -> Any:
+        if self._renderer:
+            self._renderer.render()
+
+    def close(self) -> None:
+        pass
+
+    @property
+    def num_envs(self) -> int:
+        return self._env.num_envs
+
+    @property
+    def observation_space(self) -> gymnasium.Space:
+        return self._env.observation_space
+
+    @property
+    def action_space(self) -> gymnasium.Space:
+        return self._env.action_space

motrix_rl/src/motrix_rl/skrl/torch/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+
+
+def _inherits_from(cls, base_class_name):
+    """Check if cls inherits from a class with the given base_class_name."""
+    return any(base.__name__ == base_class_name for base in cls.__mro__)
+
+
+def wrap_env(env, enable_render: bool = False):
+    """Wrap the environment based on its type."""
+    if _inherits_from(env.__class__, "NpEnv"):
+        from motrix_rl.skrl.torch.wrap_np import SkrlNpWrapper
+
+        return SkrlNpWrapper(env, enable_render=enable_render)
+    else:
+        raise ValueError(f"Unsupported environment type: {env.__class__.__name__}")

motrix_rl/src/motrix_rl/skrl/torch/train/ppo.py

@@ -0,0 +1,355 @@
+# 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 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
+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.skrl import get_log_dir
+from motrix_rl.skrl.cfg import PPOCfg
+from motrix_rl.skrl.torch import wrap_env
+
+
+def _get_cfg(
+    rlcfg: PPOCfg,
+    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
+    else:
+        cfg["rewards_shaper"] = None
+    cfg["time_limit_bootstrap"] = rlcfg.time_limit_bootstrap
+    cfg["state_preprocessor"] = RunningStandardScaler
+    cfg["state_preprocessor_kwargs"] = {
+        "size": env.observation_space,
+        "device": env.device,
+    }
+    cfg["value_preprocessor"] = RunningStandardScaler
+    cfg["value_preprocessor_kwargs"] = {"size": 1, "device": env.device}
+    # logging to TensorBoard and write checkpoints (in timesteps)
+    if log_dir:
+        cfg["experiment"]["write_interval"] = rlcfg.check_point_interval
+        cfg["experiment"]["checkpoint_interval"] = rlcfg.check_point_interval
+        cfg["experiment"]["directory"] = log_dir
+    else:
+        cfg["experiment"]["write_interval"] = 0
+        cfg["experiment"]["checkpoint_interval"] = 0
+
+    return cfg
+
+
+class PPO(BasePPO):
+    _total_custom_rewards: dict[str, torch.Tensor] = {}
+
+    def record_transition(
+        self,
+        states,
+        actions,
+        rewards,
+        next_states,
+        terminated,
+        truncated,
+        infos,
+        timestep,
+        timesteps,
+    ) -> None:
+        super().record_transition(
+            states,
+            actions,
+            rewards,
+            next_states,
+            terminated,
+            truncated,
+            infos,
+            timestep,
+            timesteps,
+        )
+        if "Reward" in infos:
+            for key, value in infos["Reward"].items():
+                value = torch.tensor(value, device=self.device)
+                self.tracking_data[f"Reward Instant / {key} (max)"].append(torch.max(value).item())
+                self.tracking_data[f"Reward Instant / {key} (min)"].append(torch.min(value).item())
+                self.tracking_data[f"Reward Instant / {key} (mean)"].append(torch.mean(value).item())
+                if key not in self._total_custom_rewards:
+                    self._total_custom_rewards[key] = torch.zeros_like(value)
+                self._total_custom_rewards[key] += value
+            done = terminated | truncated
+            done = done.reshape(-1)
+            if done.any():
+                for key in self._total_custom_rewards:
+                    self.tracking_data[f"Reward Total/ {key} (mean)"].append(
+                        torch.mean(self._total_custom_rewards[key][done]).item()
+                    )
+                    self.tracking_data[f"Reward Total/ {key} (min)"].append(
+                        torch.min(self._total_custom_rewards[key][done]).item()
+                    )
+                    self.tracking_data[f"Reward Total/ {key} (max)"].append(
+                        torch.max(self._total_custom_rewards[key][done]).item()
+                    )
+
+                    self._total_custom_rewards[key] = self._total_custom_rewards[key] * (~done)
+
+        if "metrics" in infos:
+            for key, value in infos["metrics"].items():
+                tracked_value = torch.tensor(value, device=self.device)
+                self.tracking_data[f"metrics / {key} (max)"].append(torch.max(tracked_value).item())
+                self.tracking_data[f"metrics / {key} (min)"].append(torch.min(tracked_value).item())
+                self.tracking_data[f"metrics / {key} (mean)"].append(torch.mean(tracked_value).item())
+
+
+class Trainer:
+    _trainer: SequentialTrainer
+    _env_name: str
+    _sim_backend: str
+    _rlcfg: PPOCfg
+    _enable_render: bool
+
+    def __init__(
+        self,
+        env_name: str,
+        sim_backend: str = None,
+        enable_render: bool = False,
+        cfg_override: dict = None,
+    ) -> None:
+        rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="torch")
+        if cfg_override is not None:
+            rlcfg = rlcfg.replace(**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.
+        """
+        rlcfg = self._rlcfg
+        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
+        set_seed(rlcfg.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)
+        cfg_trainer = {
+            "timesteps": rlcfg.max_batch_env_steps,
+            "headless": not self._enable_render,
+        }
+        trainer = SequentialTrainer(cfg=cfg_trainer, env=skrl_env, agents=agent)
+        trainer.train()
+
+    def play(self, policy: str) -> None:
+        import time
+
+        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)
+        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)
+        agent.load(policy)
+        with torch.no_grad():
+            obs, _ = env.reset()
+            fps = 60
+            while True:
+                t = time.time()
+                outputs = agent.act(obs, timestep=0, timesteps=0)
+                actions = outputs[-1].get("mean_actions", outputs[0])
+                obs, _, _, _, _ = env.step(actions)
+                env.render()
+                delta_time = time.time() - t
+                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."""
+            layers = []
+            current_size = input_size
+
+            for hidden_size in hidden_sizes:
+                layers.append(nn.Linear(current_size, hidden_size))
+                layers.append(activation())
+                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"]
+
+        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)
+
+        agent = PPO(
+            models=models,
+            memory=memory,
+            cfg=ppo_cfg,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            device=env.device,
+        )
+        return agent

motrix_rl/src/motrix_rl/skrl/torch/wrap_np.py

@@ -0,0 +1,79 @@
+# 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 typing import Any, Tuple
+
+import gymnasium
+import torch
+from skrl.envs.torch import Wrapper as SkrlWrapper
+
+from motrix_envs.np.env import NpEnv
+from motrix_envs.np.renderer import NpRenderer
+
+
+class SkrlNpWrapper(SkrlWrapper):
+    """
+    Wrap the numpy-based environment to be compatible with skrl (PyTorch)
+    """
+
+    _env: NpEnv
+    _renderer: NpRenderer = None
+
+    def __init__(self, env: NpEnv, enable_render: bool = False):
+        super().__init__(env)
+        if enable_render:
+            self._renderer = NpRenderer(env)
+
+    def reset(self) -> Tuple[torch.Tensor, Any]:
+        state = self._env.init_state()
+        return torch.tensor(state.obs, dtype=torch.float32, device=self.device), state.info
+
+    def step(
+        self, actions: torch.Tensor
+    ) -> Tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        Any,
+    ]:
+        actions = actions.cpu().numpy()
+        state = self._env.step(actions)
+        return (
+            torch.tensor(state.obs, dtype=torch.float32, device=self.device),
+            torch.tensor(state.reward.reshape(-1, 1), dtype=torch.float32, device=self.device),
+            torch.tensor(state.terminated.reshape(-1, 1), dtype=torch.bool, device=self.device),
+            torch.tensor(state.truncated.reshape(-1, 1), dtype=torch.bool, device=self.device),
+            state.info,
+        )
+
+    def render(self, *args, **kwargs) -> Any:
+        if self._renderer:
+            self._renderer.render()
+
+    def close(self) -> None:
+        pass
+
+    @property
+    def num_envs(self) -> int:
+        return self._env.num_envs
+
+    @property
+    def observation_space(self) -> gymnasium.Space:
+        return self._env.observation_space
+
+    @property
+    def action_space(self) -> gymnasium.Space:
+        return self._env.action_space

motrix_rl/src/motrix_rl/utils.py

@@ -0,0 +1,61 @@
+# 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
+
+
+@dataclass
+class DeviceSupports:
+    torch: bool = False
+    torch_gpu: bool = False
+    jax: bool = False
+    jax_gpu: bool = False
+
+
+def _check_gpu_available_for_torch():
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+        torch.zeros((1,)).cuda().numpy(force=True)
+        return True
+    except Exception:
+        return False
+
+
+def get_device_supports() -> DeviceSupports:
+    supports = DeviceSupports()
+    try:
+        import torch  # noqa: F401
+
+        supports.torch = True
+        supports.torch_gpu = _check_gpu_available_for_torch()
+    except ImportError:
+        pass
+
+    try:
+        import jax  # noqa: F401
+
+        supports.jax = True
+        from jax.lib import xla_bridge
+
+        platform = xla_bridge.get_backend().platform
+        if platform == "gpu":
+            supports.jax_gpu = True
+    except ImportError:
+        pass
+
+    return supports