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

chore: release v0.0.1

Browse Source
This commit is contained in:
motphys-developers
2025-11-20 08:57:48 +00:00
commit 5133830b5a
105 changed files with 8789 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

172
motrix_envs/README.md Normal file
View File

@@ -0,0 +1,172 @@
# Motrix Environments (Motrix_envs)
motrix Environments 是一个基于 MotrixSim 仿真后端的强化学习环境库，提供了与具体 RL 框架无关的仿真环境定义。该模块设计上支持多种仿真后端，目前主要实现了 MotrixSim 的 NumPy 后端。
## 📁 项目结构
```
motrix_envs/
├── base.py # 基础抽象类和配置
├── registry.py # 环境注册系统
├── np/ # NumPy 仿真后端实现
│ ├── env.py # NumPy 环境基类
│ ├── renderer.py # 渲染器
│ └── reward.py # 奖励函数
├── basic/ # 基础环境
│ ├── cartpole/ # 倒立摆环境
│ └── walker/ # 步行者环境
├── locomotion/ # 运动控制环境
│ └── go1/ # GO1 机器人
│ ├── xmls/ # 机器人模型文件
│ ├── walk_np.py # GO1 行走实现
│ └── cfg.py # GO1 配置
└── common/ # 公共组件
```
## 🎯 内置环境
| 环境名称 | 注册标识符 | 后端 | 类型 | 描述 |
| ---------------- | ----------------------- | ---- | ---------- | ----------------------------- |
| **倒立摆** | `cartpole` | np | Basic | 经典控制任务，保持杆子平衡 |
| **步行者** | `walker` | np | Basic | 平面双足步行机器人控制 |
| **GO1 平地行走** | `go1-flat-terrain-walk` | np | Locomotion | 四足机器人 GO1 的平地行走任务 |
### 详细说明
#### 1. CartPole (倒立摆) - `cartpole`
- **观测空间 (4 维)**:
- `cart_pos`: 小车位置 [-0.8, 0.8]
- `pole_angle`: 杆子角度 [-0.2, 0.2]
- `cart_vel`: 小车速度
- `pole_vel`: 杆子角速度
- **动作空间 (1 维)**: 推力 `[-3.0, 3.0]`
- **奖励**: 每步 +1.0
- **终止条件**: 杆子角度 > 0.2 弧度 或 小车位置超出边界
- **配置参数**: `reset_noise_scale=0.01`
#### 2. GO1 平地行走 - `go1-flat-terrain-walk`
- **观测空间 (48 维)**: 包含关节位置、速度、IMU 数据、命令等
- **动作空间 (12 维)**: 12 个关节的目标角度控制
- **奖励组件**:
- `tracking_lin_vel`: 线速度跟踪 (权重: 1.0)
- `tracking_ang_vel`: 角速度跟踪 (权重: 0.5)
- `feet_air_time`: 足部空中时间 (权重: 1.0)
- `collision`: 碰撞惩罚 (权重: -1.0)
- `action_rate`: 动作变化率惩罚 (权重: -0.001)
- **控制参数**: 刚度 80.0 Nm/rad, 阻尼 1.0 Nms/rad
- **噪声模型**: 关节角度、速度、陀螺仪等多种传感器噪声
## 🛠️ 自定义环境
### 开发步骤
#### 1. 定义环境配置类
```python
from dataclasses import dataclass
from motrix_envs import registry
from motrix_envs.base import EnvCfg
@registry.envcfg("my-custom-env")
@dataclass
class MyEnvCfg(EnvCfg):
"""自定义环境配置"""
# 继承基础配置
reset_noise_scale: float = 0.01
max_episode_seconds: float = 10.0
# 添加自定义参数
custom_param_1: float = 1.0
custom_param_2: str = "default_value"
enable_feature_x: bool = True
```
#### 2. 实现环境类
```python
import gymnasium as gym
import numpy as np
from motrix_envs import registry
from motrix_envs.np.env import NpEnv, NpEnvState
@registry.env("my-custom-env", "np")
class MyCustomEnv(NpEnv):
"""自定义环境实现"""
def __init__(self, cfg: MyEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
@property
def observation_space(self):
"""返回观测空间"""
raise NotImplementedError
@property
def action_space(self):
"""返回动作空间"""
raise NotImplementedError
def apply_action(self, actions: np.ndarray, state: NpEnvState):
"""将actions应用到环境状态"""
raise NotImplementedError
return state
def update_state(self, state: NpEnvState):
"""更新环境状态，计算观测、奖励、终止条件"""
# 提取仿真数据
data = state.data
# 计算观测
obs = self._compute_observation(data)
# 计算奖励
reward = self._compute_reward(data, obs)
# 检查终止条件
terminated = self._check_termination(data, obs)
# 更新状态
state.obs = obs
state.reward = reward
state.terminated = terminated
return state
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
"""
重置环境
参数:
data: motrixsim 场景数据对象，仅包含需要重置的环境实例
返回:
obs: 重置后场景的初始观测
info: 额外信息
"""
raise NotImplementedError
```
#### 3. 注册和使用环境
```python
# 确保导入了环境模块，触发注册
import motrix_envs.basic.my_custom_env # noqa: F401
from motrix_envs import registry
# 创建环境实例
env = registry.make(
name="my-custom-env",
sim_backend="np",
num_envs=256,
env_cfg_override={
"custom_param_1": 2.0,
"reset_noise_scale": 0.02
}
)
# 使用环境
for step in range(1000):
actions = sample_actions_somehow()
state = env.step(actions)
```

15
motrix_envs/pyproject.toml Normal file
View File

@@ -0,0 +1,15 @@
[build-system]
requires = ["uv_build>=0.9.6,<0.10.0"]
build-backend = "uv_build"
[project]
name = "motrix-envs"
version = "0.0.1"
description = "Robot simulation environment library based on MotrixSim providing multi-task RL environments."
authors = [{ name = "Motphys", email = "developers@motphys.com" }]
requires-python = "==3.10.*"
readme = "README.md"
license = "Apache-2.0"
dependencies = [
"motrixsim>=0.4.0",
]

16
motrix_envs/src/motrix_envs/__init__.py Normal file
View File

@@ -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 basic, locomotion # noqa: F401

82
motrix_envs/src/motrix_envs/base.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.
# ==============================================================================
import abc
from dataclasses import dataclass
from typing import Optional
import gymnasium as gym
@dataclass
class EnvCfg:
"""
Config for the environment
"""
model_file: str = None
sim_dt: float = 0.01
max_episode_seconds: float = None
ctrl_dt: float = 0.01
@property
def max_episode_steps(self) -> Optional[int]:
"""
return the max episode steps
"""
if self.max_episode_seconds is None:
return None
return int(self.max_episode_seconds / self.ctrl_dt)
@property
def sim_substeps(self) -> int:
"""
return the number of simulation steps per control step
"""
return int(round(self.ctrl_dt / self.sim_dt))
def validate(self):
"""
validate the config
"""
if self.sim_dt > self.ctrl_dt:
raise ValueError("sim_dt must be less than or equal to ctrl_dt")
class ABEnv(abc.ABC):
@property
@abc.abstractmethod
def num_envs(self) -> int:
"""
return the size of the env if it is vectorized
"""
@property
@abc.abstractmethod
def cfg(self) -> EnvCfg:
"""
The configuration of the environment
"""
@property
@abc.abstractmethod
def observation_space(self) -> gym.Space:
"""Observation space"""
@property
@abc.abstractmethod
def action_space(self) -> gym.Space:
"""Action space"""

16
motrix_envs/src/motrix_envs/basic/__init__.py Normal file
View File

@@ -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 cartpole, walker # noqa: F401 import to register envs

16
motrix_envs/src/motrix_envs/basic/cartpole/__init__.py Normal file
View File

@@ -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 cartpole_np # noqa: F401

36
motrix_envs/src/motrix_envs/basic/cartpole/cartpole.xml Normal file
View File

@@ -0,0 +1,36 @@
<mujoco model="cartpole">
<compiler inertiafromgeom="true" />
<option gravity="0 0 -9.81" integrator="RK4" timestep="0.01" />
<default>
<joint armature="0" damping="1" limited="true" solimplimit="0 0 0.001 0.5 1.0" />
<geom contype="0" conaffinity="0" friction="1 0.1 0.1" />
<motor ctrlrange="-3 3" />
</default>
<asset>
<texture type="2d" name="groundplane" builtin="checker" mark="edge"
rgb1="0.4314 0.5294 0.6431" rgb2="0.8157 0.8549 0.9059"
markrgb="0.8 0.8 0.8" width="300" height="300" />
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="1 1"
reflectance="0.2" />
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true" />
<geom name="floor" size="0 0 0.01" type="plane" material="groundplane" pos="0 0 -1" />
</worldbody>
<worldbody>
<geom name="rail" pos="0 0 0" quat="0.707 0 0.707 0" rgba="0.3 0.3 0.7 1" size="0.02 1"
type="capsule" />
<body name="cart" pos="0 0 0">
<joint axis="1 0 0" limited="true" name="slider" pos="0 0 0" range="-1 1" type="slide" />
<geom name="cart" pos="0 0 0" quat="0.707 0 0.707 0" size="0.1 0.1" type="capsule" />
<body name="pole" pos="0 0 0">
<joint axis="0 1 0" name="hinge" pos="0 0 0" range="-90 90" type="hinge" />
<geom pos="0 0 0.3" name="cpole" rgba="0 0.7 0.7 1" size="0.049 0.3" type="capsule" />
</body>
</body>
</worldbody>
<actuator>
<motor ctrllimited="true" ctrlrange="-3 3" gear="100" joint="slider" name="slide" />
</actuator>
</mujoco>

97
motrix_envs/src/motrix_envs/basic/cartpole/cartpole_np.py Normal file
View File

@@ -0,0 +1,97 @@
# 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 gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import CartPoleEnvCfg
@registry.env("cartpole", "np")
class CartPoleEnv(NpEnv):
_cfg: CartPoleEnvCfg
def __init__(self, cfg: CartPoleEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
self._action_space = gym.spaces.Box(-3.0, 3.0, (1,), dtype=np.float32)
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (4,), dtype=np.float32)
self._num_dof_pos = self._model.num_dof_pos
self._num_dof_vel = self._model.num_dof_vel
self._init_dof_pos = self._model.compute_init_dof_pos()
self._init_dof_vel = np.zeros(
(self._model.num_dof_vel,),
dtype=np.float32,
)
@property
def observation_space(self):
return self._observation_space
@property
def action_space(self):
return self._action_space
def apply_action(self, actions: np.ndarray, state: NpEnvState):
state.data.actuator_ctrls = actions
return state
def update_state(self, state: NpEnvState):
# compute observation
data = state.data
dof_pos = data.dof_pos
dof_vel = data.dof_vel
obs = np.concatenate([dof_pos, dof_vel], axis=-1)
assert obs.shape == (self._num_envs, 4)
# compute reward
reward = np.ones((self._num_envs,), dtype=np.float32)
# compute terminated
cart_pos = dof_pos[:, 0]
angle = dof_pos[:, 1]
terminated = np.logical_or(np.isnan(angle), np.abs(angle) > 0.2)
terminated = np.logical_or(cart_pos < -0.8, terminated)
terminated = np.logical_or(cart_pos > 0.8, terminated)
state.obs = obs
state.reward = reward
state.terminated = terminated
return state
def reset(self, data: mtx.SceneData):
cfg: CartPoleEnvCfg = self._cfg
noise_pos = np.random.uniform(
-cfg.reset_noise_scale,
cfg.reset_noise_scale,
(*data.shape, self._num_dof_pos),
)
noise_vel = np.random.uniform(
-cfg.reset_noise_scale,
cfg.reset_noise_scale,
(*data.shape, self._num_dof_vel),
)
dof_pos = np.tile(self._init_dof_pos, (*data.shape, 1)) + noise_pos
dof_vel = np.tile(self._init_dof_vel, (*data.shape, 1)) + noise_vel
data.reset(self._model)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
obs = np.concatenate([dof_pos, dof_vel], axis=-1)
return obs, {}

30
motrix_envs/src/motrix_envs/basic/cartpole/cfg.py Normal file
View File

@@ -0,0 +1,30 @@
# 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 os
from dataclasses import dataclass
from motrix_envs import registry
from motrix_envs.base import EnvCfg
model_file = os.path.dirname(__file__) + "/cartpole.xml"
@registry.envcfg("cartpole")
@dataclass
class CartPoleEnvCfg(EnvCfg):
model_file: str = model_file
reset_noise_scale: float = 0.01
max_episode_seconds: float = 10

16
motrix_envs/src/motrix_envs/basic/walker/__init__.py Normal file
View File

@@ -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 walker_np # noqa: F401

45
motrix_envs/src/motrix_envs/basic/walker/cfg.py Normal file
View File

@@ -0,0 +1,45 @@
# 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 os
from dataclasses import dataclass
from motrix_envs import registry
from motrix_envs.base import EnvCfg
model_file = os.path.dirname(__file__) + "/walker.xml"
@registry.envcfg("dm-walker")
@dataclass
class WalkerEnvCfg(EnvCfg):
model_file: str = model_file
max_episode_seconds: float = 25.0
sim_dt: float = 0.0125
move_speed: float = 1.0
ctrl_dt: float = 0.025
stand_height: float = 1.2
@registry.envcfg("dm-stander")
@dataclass
class StanderEnvCfg(WalkerEnvCfg):
move_speed: float = 0.0
@registry.envcfg("dm-runner")
@dataclass
class RunnerEnvCfg(WalkerEnvCfg):
move_speed: float = 5.0

71
motrix_envs/src/motrix_envs/basic/walker/walker.xml Normal file
View File

@@ -0,0 +1,71 @@
<mujoco model="planar walker">
<include file="../../common/visual.xml" />
<include file="../../common/skybox.xml" />
<include file="../../common/materials.xml" />
<option timestep="0.005">
<flag eulerdamp="disable" />
</option>
<statistic extent="2" center="0 0 1" />
<default>
<joint damping=".1" armature="0.01" limited="true" solimplimit="0 .99 .01" />
<geom contype="1" conaffinity="1" friction=".7 0.01 0.01" />
<!-- Note: the original friction is 0.1 -->
<motor ctrlrange="-1 1" ctrllimited="true" />
<site size="0.01" />
<default class="walker">
<geom material="self" type="capsule" />
<joint axis="0 -1 0" />
</default>
</default>
<worldbody>
<light name="light" dir="0 -1 -1" directional="true"/>
<geom name="floor" type="plane" conaffinity="1" pos="248 0 0" size="250 .8 .2" material="grid" zaxis="0 0 1" />
<body name="torso" pos="0 0 1.3" childclass="walker">
<joint name="rootz" axis="0 0 1" type="slide" limited="false" armature="0" damping="0" />
<joint name="rootx" axis="1 0 0" type="slide" limited="false" armature="0" damping="0" />
<joint name="rooty" axis="0 1 0" type="hinge" limited="false" armature="0" damping="0" />
<geom name="torso" size="0.07 0.3" />
<body name="right_thigh" pos="0 -.05 -0.3">
<joint name="right_hip" range="-20 100" />
<geom name="right_thigh" pos="0 0 -0.225" size="0.05 0.225" />
<body name="right_leg" pos="0 0 -0.7">
<joint name="right_knee" pos="0 0 0.25" range="-150 0" />
<geom name="right_leg" size="0.04 0.25" />
<body name="right_foot" pos="0.06 0 -0.25">
<joint name="right_ankle" pos="-0.06 0 0" range="-45 45" />
<geom name="right_foot" zaxis="1 0 0" size="0.05 0.1" />
</body>
</body>
</body>
<body name="left_thigh" pos="0 .05 -0.3">
<joint name="left_hip" range="-20 100" />
<geom name="left_thigh" pos="0 0 -0.225" size="0.05 0.225" />
<body name="left_leg" pos="0 0 -0.7">
<joint name="left_knee" pos="0 0 0.25" range="-150 0" />
<geom name="left_leg" size="0.04 0.25" />
<body name="left_foot" pos="0.06 0 -0.25">
<joint name="left_ankle" pos="-0.06 0 0" range="-45 45" />
<geom name="left_foot" zaxis="1 0 0" size="0.05 0.1" />
</body>
</body>
</body>
</body>
</worldbody>
<sensor>
<subtreelinvel name="torso_subtreelinvel" body="torso" />
</sensor>
<actuator>
<motor name="right_hip" joint="right_hip" gear="100" />
<motor name="right_knee" joint="right_knee" gear="50" />
<motor name="right_ankle" joint="right_ankle" gear="20" />
<motor name="left_hip" joint="left_hip" gear="100" />
<motor name="left_knee" joint="left_knee" gear="50" />
<motor name="left_ankle" joint="left_ankle" gear="20" />
</actuator>
</mujoco>

155
motrix_envs/src/motrix_envs/basic/walker/walker_np.py Normal file
View File

@@ -0,0 +1,155 @@
# 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 gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.basic.walker.cfg import WalkerEnvCfg
from motrix_envs.np import reward
from motrix_envs.np.env import NpEnv, NpEnvState
@registry.env("dm-walker", "np")
@registry.env("dm-runner", "np")
@registry.env("dm-stander", "np")
class Walker2DEnv(NpEnv):
_observation_space: gym.spaces.Box
_action_space: gym.spaces.Box
def __init__(self, cfg: WalkerEnvCfg, num_envs=1):
super().__init__(cfg, num_envs)
self._init_obs_space()
self._init_action_space()
self._torso = self._model.get_link("torso")
self._move_speed = cfg.move_speed
self._joint_limits = self._model.joint_limits
self._stand_height = cfg.stand_height
def _init_obs_space(self):
model = self._model
num = 0
num += (model.num_links - 1) * 2 # planar rotation (x,z) for each link except the root
num += 1 # torso height
num += model.num_dof_vel
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (num,), dtype=np.float32)
def _init_action_space(self):
model = self._model
self._action_space = gym.spaces.Box(
model.actuator_ctrl_limits[0],
model.actuator_ctrl_limits[1],
(model.num_actuators,),
dtype=np.float32,
)
@property
def observation_space(self) -> gym.spaces.Box:
return self._observation_space
@property
def action_space(self) -> gym.spaces.Box:
return self._action_space
def apply_action(self, actions, state):
state.data.actuator_ctrls = actions
return state
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
num_env = data.shape[0]
link_rotations = self._model.get_link_rotation_mats(data)
dof_vel = data.dof_vel
up_right = link_rotations[:, 0, 2, 2].reshape(num_env, 1) # 1
orientations = link_rotations[:, 1:, [0, 0], [2, 0]].reshape(num_env, -1) # (num_links - 1) * 2
obs = np.concatenate([orientations, up_right, dof_vel], axis=-1)
return obs
def update_state(self, state: NpEnvState) -> NpEnvState:
data = state.data
# === compute obs ====
obs = self._get_obs(data)
offset = (self._model.num_links - 1) * 2 + 1
dof_vel = obs[:, offset:]
torso_upright = self._torso.get_rotation_mat(data)[:, 2, 2]
torso_height = self._torso.get_position(data)[:, 2]
torso_vel = self._model.get_sensor_value("torso_subtreelinvel", data)
horizontal_velocity = torso_vel[:, 0]
# ==== compute terminated
terminated = np.isnan(dof_vel).any(axis=-1)
# ==== compute reward
rwd_height = reward.tolerance(
torso_height,
bounds=(self._stand_height, float("inf")),
margin=self._stand_height * 4 / 5,
)
rwd_upright = (1 + torso_upright) / 2
rwd_stand = (3 * rwd_height + 1 * rwd_upright) / 4
rwd = rwd_stand
state.info["Reward"] = {
"height": rwd_height,
"upright": rwd_upright,
"stand": rwd_stand,
}
if self._move_speed > 0.0:
rwd_move = reward.tolerance(
horizontal_velocity,
bounds=(self._move_speed, float("inf")),
margin=self._move_speed / 2,
value_at_margin=0.5,
sigmoid="linear",
)
state.info["Reward"]["move"] = rwd_move
rwd = rwd_stand * (5 * rwd_move + 1) / 6
rwd[terminated] = 0.0
return state.replace(
obs=obs,
reward=rwd,
terminated=terminated,
)
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
data.reset(self._model)
num_reset = data.shape[0]
dof_pos = np.zeros((num_reset, self._model.num_dof_pos))
dof_pos[:, 2] = np.random.uniform(low=-np.pi, high=np.pi, size=(num_reset,)) # randomize root yaw
dof_pos[:, 3:] = np.random.uniform(
low=self._joint_limits[0, 3:],
high=self._joint_limits[1, 3:],
size=(num_reset, self._model.num_dof_pos - 3),
) # randomize other joint angles
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
obs = self._get_obs(data)
rewards = {
"height": np.zeros((num_reset,)),
"upright": np.zeros((num_reset,)),
"stand": np.zeros((num_reset,)),
}
if self._move_speed > 0.0:
rewards["move"] = np.zeros((num_reset,))
return obs, {"Reward": rewards}

23
motrix_envs/src/motrix_envs/common/materials.xml Normal file
View File

@@ -0,0 +1,23 @@
<!--
Common textures, colors and materials to be used throughout this suite. Some
materials such as xxx_highlight are activated on occurence of certain events,
for example receiving a positive reward.
-->
<mujoco>
<asset>
<texture name="grid" type="2d" builtin="checker" mark="edge" rgb1="0.4314 0.5294 0.6431" rgb2="0.8157 0.8549 0.9059" markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
<material name="self" rgba=".7 .5 .3 1"/>
<material name="self_default" rgba=".7 .5 .3 1"/>
<material name="self_highlight" rgba="0 .5 .3 1"/>
<material name="effector" rgba=".7 .4 .2 1"/>
<material name="effector_default" rgba=".7 .4 .2 1"/>
<material name="effector_highlight" rgba="0 .5 .3 1"/>
<material name="decoration" rgba=".3 .5 .7 1"/>
<material name="eye" rgba="0 .2 1 1"/>
<material name="target" rgba=".6 .3 .3 1"/>
<material name="target_default" rgba=".6 .3 .3 1"/>
<material name="target_highlight" rgba=".6 .3 .3 .4"/>
<material name="site" rgba=".5 .5 .5 .3"/>
</asset>
</mujoco>

6
motrix_envs/src/motrix_envs/common/skybox.xml Normal file
View File

@@ -0,0 +1,6 @@
<mujoco>
<asset>
<texture name="skybox" type="skybox" builtin="gradient" rgb1=".4 .6 .8" rgb2="0 0 0"
width="800" height="800" mark="random" markrgb="1 1 1"/>
</asset>
</mujoco>

7
motrix_envs/src/motrix_envs/common/visual.xml Normal file
View File

@@ -0,0 +1,7 @@
<mujoco>
<visual>
<headlight ambient=".4 .4 .4" diffuse=".8 .8 .8" specular="0.1 0.1 0.1"/>
<map znear=".01"/>
<quality shadowsize="2048"/>
</visual>
</mujoco>

16
motrix_envs/src/motrix_envs/locomotion/__init__.py Normal file
View File

@@ -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 go1 # noqa: F401 register envs

16
motrix_envs/src/motrix_envs/locomotion/go1/__init__.py Normal file
View File

@@ -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 walk_np # noqa: F401 register envs

140
motrix_envs/src/motrix_envs/locomotion/go1/cfg.py Normal file
View File

@@ -0,0 +1,140 @@
# 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 os
from dataclasses import dataclass, field
from motrix_envs import registry
from motrix_envs.base import EnvCfg
model_file = os.path.dirname(__file__) + "/xmls/scene_motor_actuator.xml"
@dataclass
class NoiseConfig:
level: float = 1.0
scale_joint_angle: float = 0.03
scale_joint_vel: float = 1.5
scale_gyro: float = 0.2
scale_gravity: float = 0.05
scale_linvel: float = 0.1
@dataclass
class ControlConfig:
stiffness = 80 # [N*m/rad]
damping = 1 # [N*m*s/rad]
# action scale: target angle = actionScale * action + defaultAngle
action_scale = 0.05
@dataclass
class InitState:
# the initial position of the robot in the world frame
pos = [0.0, 0.0, 0.42]
# the default angles for all joints. key = joint name, value = target angle [rad]
default_joint_angles = {
"FL_hip": 0.0, # [rad]
"RL_hip": 0.0, # [rad]
"FR_hip": -0.0, # [rad]
"RR_hip": -0.0, # [rad]
"FL_thigh": 0.9, # [rad]
"RL_thigh": 0.9, # [rad]
"FR_thigh": 0.9, # [rad]
"RR_thigh": 0.9, # [rad]
"FL_calf": -1.8, # [rad]
"RL_calf": -1.8, # [rad]
"FR_calf": -1.8, # [rad]
"RR_calf": -1.8, # [rad]
}
@dataclass
class Commands:
vel_limit = [
[0.0, -1.0, -1.0], # min: vel_x [m/s], vel_y [m/s], ang_vel [rad/s]
[2.0, 1.0, 1.0], # max
]
@dataclass
class Normalization:
lin_vel = 2
ang_vel = 0.25
dof_pos = 1
dof_vel = 0.05
@dataclass
class Asset:
body_name = "trunk"
foot_name = "foot"
penalize_contacts_on = ["thigh", "calf"]
terminate_after_contacts_on = [
"trunk",
]
ground = "floor"
@dataclass
class Sensor:
local_linvel = "local_linvel"
gyro = "gyro"
@dataclass
class RewardConfig:
scales: dict[str, float] = field(
default_factory=lambda: {
"termination": -0.0,
"tracking_lin_vel": 1.0,
"tracking_ang_vel": 0.5,
"lin_vel_z": -2.0,
"ang_vel_xy": -0.05,
"orientation": -0.0,
"torques": -0.00001,
"dof_vel": -0.0,
"dof_acc": -2.5e-7,
"base_height": -0.0,
"feet_air_time": 1.0,
"collision": -1.0 * 0,
"feet_stumble": -0.0,
"action_rate": -0.001,
"stand_still": -0.0,
"hip_pos": -1,
"calf_pos": -0.3 * 0,
}
)
tracking_sigma: float = 0.25
max_foot_height: float = 0.1
@registry.envcfg("go1-flat-terrain-walk")
@dataclass
class Go1WalkNpEnvCfg(EnvCfg):
max_episode_seconds: float = 20.0
model_file: str = model_file
noise_config: NoiseConfig = field(default_factory=NoiseConfig)
control_config: ControlConfig = field(default_factory=ControlConfig)
reward_config: RewardConfig = field(default_factory=RewardConfig)
init_state: InitState = field(default_factory=InitState)
commands: Commands = field(default_factory=Commands)
normalization: Normalization = field(default_factory=Normalization)
asset: Asset = field(default_factory=Asset)
sensor: Sensor = field(default_factory=Sensor)
sim_dt: float = 0.01
ctrl_dt: float = 0.01

413
motrix_envs/src/motrix_envs/locomotion/go1/walk_np.py Normal file
View File

@@ -0,0 +1,413 @@
# 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 gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpEnvCfg
from motrix_envs.np.env import NpEnv, NpEnvState
## provide quat math utility from motrixsim.
def quat_rotate_inverse(quats, v):
"""
Rotate a fixed vector v by a list of quaternions using a vectorized approach.
Parameters:
quats (np.ndarray): Array of quaternions of shape (N, 4). Each quaternion is in [w, x, y, z] format.
v (np.ndarray): Fixed vector of shape (3,) to be rotated.
Returns:
np.ndarray: Array of rotated vectors of shape (N, 3).
"""
# Normalize the quaternions to ensure they are unit quaternions
# Extract the scalar (w) and vector (x, y, z) parts of the quaternions
w = quats[:, -1] # Shape (N,)
im = quats[:, :3] # Shape (N, 3)
# Compute the cross product between the imaginary part of each quaternion and the fixed vector v.
# np.cross broadcasts v to match each row in im, resulting in an array of shape (N, 3)
cross_im_v = np.cross(im, v)
# Compute the intermediate terms for the rotation formula:
term1 = w[:, np.newaxis] * cross_im_v # w * cross(im, v)
term2 = np.cross(im, cross_im_v) # cross(im, cross(im, v))
# Apply the rotation formula: v_rot = v + 2 * (term1 + term2)
v_rotated = v + 2 * (term1 + term2)
return v_rotated
@registry.env("go1-flat-terrain-walk", sim_backend="np")
class Go1WalkTask(NpEnv):
_init_dof_pos: np.ndarray
_init_dof_vel: np.ndarray
def __init__(self, cfg: Go1WalkNpEnvCfg, num_envs=1):
super().__init__(cfg, num_envs)
self._init_action_space()
self._init_obs_space()
self._body = self._model.get_body(self.cfg.asset.body_name)
self._num_action = self._action_space.shape[0]
self._num_observation = self._observation_space.shape[0]
self._num_dof_pos = self._model.num_dof_pos
self._num_dof_vel = self._model.num_dof_vel
self._init_dof_vel = np.zeros(
(self._num_dof_vel,),
dtype=np.float32,
)
self._init_dof_pos = self._model.compute_init_dof_pos()
self._init_buffer()
def _init_obs_space(self):
model = self.model
num_dof_vel = model.num_dof_vel # linvel + gyro + joint_vel
num_joint_angle = model.num_dof_pos - 7
num_gravity = 3
num_actions = model.num_actuators
num_command = 3
num_obs = num_dof_vel + num_joint_angle + num_gravity + num_actions + num_command
assert num_obs == 48
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (num_obs,), dtype=np.float32)
def _init_action_space(self):
model = self.model
self._action_space = gym.spaces.Box(
np.array(model.actuator_ctrl_limits[0, :]),
np.array(model.actuator_ctrl_limits[1, :]),
(model.num_actuators,),
dtype=np.float32,
)
@property
def action_space(self) -> gym.spaces.Box:
return self._action_space
@property
def observation_space(self) -> gym.spaces.Box:
return self._observation_space
def get_dof_pos(self, data: mtx.SceneModel):
return self._body.get_joint_dof_pos(data)
def get_dof_vel(self, data: mtx.SceneModel):
return self._body.get_joint_dof_vel(data)
def _init_buffer(self):
cfg = self._cfg
assert isinstance(cfg, Go1WalkNpEnvCfg)
# init buffers
self.reset_buf = np.ones(self._num_envs, dtype=np.bool)
self.kps = np.ones(self._num_action, dtype=np.float32) * cfg.control_config.stiffness
self.kds = np.ones(self._num_action, dtype=np.float32) * cfg.control_config.damping
self.gravity_vec = np.array([0, 0, -1], dtype=np.float32)
self.commands_scale = np.array(
(
[
cfg.normalization.lin_vel,
cfg.normalization.lin_vel,
cfg.normalization.ang_vel,
]
),
dtype=np.float32,
)
self.default_angles = np.zeros(self._num_action, dtype=np.float32)
self.hip_indices = []
self.calf_indices = []
for i in range(self._model.num_actuators):
for name in cfg.init_state.default_joint_angles.keys():
if name in self._model.actuator_names[i]:
self.default_angles[i] = cfg.init_state.default_joint_angles[name]
if "hip" in self._model.actuator_names[i]:
self.hip_indices.append(i)
if "calf" in self._model.actuator_names[i]:
self.calf_indices.append(i)
self._init_dof_pos[-self._num_action :] = self.default_angles
self.ground = self._model.get_geom_index(cfg.asset.ground)
self.termination_contact = None
self.foot = []
for name in cfg.asset.terminate_after_contacts_on:
if self.termination_contact is None:
self.termination_contact = np.array([[self._model.get_geom_index(name), self.ground]], dtype=np.uint32)
else:
self.termination_contact = np.append(
self.termination_contact,
np.array(
[[self._model.get_geom_index(name), self.ground]],
dtype=np.uint32,
),
axis=0,
)
for name in cfg.asset.foot_name:
self.foot.append([self._model.get_geom_index(name), self.ground])
self.num_check = self.termination_contact.shape[0]
self.foot = None
for i in self._model.geom_names:
if i is not None and cfg.asset.foot_name in i:
if self.foot is None:
self.foot = np.array([[self._model.get_geom_index(i), self.ground]], dtype=np.uint32)
else:
self.foot = np.append(
self.foot,
np.array(
[[self._model.get_geom_index(i), self.ground]],
dtype=np.uint32,
),
axis=0,
)
self.foot_check_num = self.foot.shape[0]
self.foot_check = self.foot
self.termination_check = self.termination_contact
def apply_action(self, actions, state):
state.info["last_dof_vel"] = self.get_dof_vel(state.data)
state.info["last_actions"] = state.info["current_actions"]
state.info["current_actions"] = actions
state.data.actuator_ctrls = self._compute_torques(actions, state.data)
return state
def _compute_torques(self, actions, data):
# Compute torques from actions.
# pd controller
actions_scaled = actions * self.cfg.control_config.action_scale
torques = self.kps * (
actions_scaled + self.default_angles - self.get_dof_pos(data)
) - self.kds * self.get_dof_vel(data)
return torques
def get_local_linvel(self, data: mtx.SceneData) -> np.ndarray:
return self._model.get_sensor_value(self.cfg.sensor.local_linvel, data)
def get_gyro(self, data: mtx.SceneData) -> np.ndarray:
return self._model.get_sensor_value(self.cfg.sensor.gyro, data)
def update_state(self, state):
state = self.update_observation(state)
state = self.update_terminated(state)
state = self.update_reward(state)
return state
def _get_obs(self, data: mtx.SceneData, info: dict) -> np.ndarray:
linear_vel = self.get_local_linvel(data)
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
local_gravity = quat_rotate_inverse(base_quat, self.gravity_vec)
diff = self.get_dof_pos(data) - self.default_angles
noisy_linvel = linear_vel * self.cfg.normalization.lin_vel
noisy_gyro = gyro * self.cfg.normalization.ang_vel
noisy_joint_angle = diff * self.cfg.normalization.dof_pos
noisy_joint_vel = self.get_dof_vel(data) * self.cfg.normalization.dof_vel
command = info["commands"] * self.commands_scale
last_actions = info["current_actions"]
obs = np.hstack(
[
noisy_linvel,
noisy_gyro,
local_gravity,
noisy_joint_angle,
noisy_joint_vel,
last_actions,
command,
]
)
return obs
def update_observation(self, state: NpEnvState):
data = state.data
obs = self._get_obs(data, state.info)
cquerys = self._model.get_contact_query(data)
foot_contact = cquerys.is_colliding(self.foot_check)
state.info["contacts"] = foot_contact.reshape((self._num_envs, self.foot_check_num))
state.info["feet_air_time"] = self.update_feet_air_time(state.info)
return state.replace(obs=obs)
def update_terminated(self, state: NpEnvState) -> NpEnvState:
data = state.data
cquerys = self._model.get_contact_query(data)
termination_check = cquerys.is_colliding(self.termination_check)
termination_check.reshape((self._num_envs, self.num_check))
terminated = termination_check.any(axis=1)
return state.replace(
terminated=terminated,
)
def update_feet_air_time(self, info: dict):
feet_air_time = info["feet_air_time"]
feet_air_time += self.cfg.ctrl_dt
feet_air_time *= ~info["contacts"]
return feet_air_time
def resample_commands(self, num_envs: int):
commands = np.random.uniform(
low=self.cfg.commands.vel_limit[0],
high=self.cfg.commands.vel_limit[1],
size=(num_envs, 3),
)
return commands
def update_reward(self, state: NpEnvState) -> NpEnvState:
data = state.data
terminated = state.terminated
reward_dict = self._get_reward(data, state.info)
rewards = {k: v * self.cfg.reward_config.scales[k] for k, v in reward_dict.items()}
rwd = sum(rewards.values())
rwd = np.clip(rwd, 0.0, 10000.0)
if "termination" in self.cfg.reward_config.scales:
termination = self._reward_termination(terminated) * self.cfg.reward_config.scales["termination"]
rwd += termination
rwd = np.where(terminated, np.array(0.0), rwd)
return state.replace(reward=rwd)
def reset(self, data) -> tuple[np.ndarray, dict]:
num_reset = data.shape[0]
dof_pos = np.tile(self._init_dof_pos, (num_reset, 1))
dof_vel = np.tile(self._init_dof_vel, (num_reset, 1))
data.reset(self._model)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
info = {
"current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"commands": self.resample_commands(num_reset),
"last_dof_vel": np.zeros((num_reset, self._num_action), dtype=np.float32),
"feet_air_time": np.zeros((num_reset, self.foot_check_num), dtype=np.float32),
"contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
}
obs = self._get_obs(data, info)
return obs, info
def _get_reward(
self,
data: mtx.SceneData,
info: dict,
) -> dict[str, np.ndarray]:
commands = info["commands"]
return {
"lin_vel_z": self._reward_lin_vel_z(data),
"ang_vel_xy": self._reward_ang_vel_xy(data),
"orientation": self._reward_orientation(data),
"torques": self._reward_torques(data),
"dof_vel": self._reward_dof_vel(data),
"dof_acc": self._reward_dof_acc(data, info),
"action_rate": self._reward_action_rate(info),
"tracking_lin_vel": self._reward_tracking_lin_vel(data, commands),
"tracking_ang_vel": self._reward_tracking_ang_vel(data, commands),
"stand_still": self._reward_stand_still(data, commands),
"hip_pos": self._reward_hip_pos(data, commands),
"calf_pos": self._reward_calf_pos(data, commands),
"feet_air_time": self._reward_feet_air_time(commands, info),
}
# ------------ reward functions----------------
def _reward_lin_vel_z(self, data):
# Penalize z axis base linear velocity
return np.square(self.get_local_linvel(data)[:, 2])
def _reward_ang_vel_xy(self, data):
# Penalize xy axes base angular velocity
return np.sum(np.square(self.get_gyro(data)[:, :2]), axis=1)
def _reward_orientation(self, data):
# Penalize non flat base orientation
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = quat_rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):
# Penalize torques
return np.sum(np.square(data.actuator_ctrls), axis=1)
def _reward_dof_vel(self, data):
# Penalize dof velocities
return np.sum(np.square(self.get_dof_vel(data)), axis=1)
def _reward_dof_acc(self, data, info):
# Penalize dof accelerations
return np.sum(
np.square((info["last_dof_vel"] - self.get_dof_vel(data)) / self.cfg.ctrl_dt),
axis=1,
)
def _reward_action_rate(self, info: dict):
# Penalize changes in actions
action_diff = info["current_actions"] - info["last_actions"]
return np.sum(np.square(action_diff), axis=1)
def _reward_termination(self, done):
# Terminal reward / penalty
return done
def _reward_feet_air_time(self, commands: np.ndarray, info: dict):
# Reward long steps
feet_air_time = info["feet_air_time"]
first_contact = (feet_air_time > 0.0) * info["contacts"]
# reward only on first contact with the ground
rew_airTime = np.sum((feet_air_time - 0.5) * first_contact, axis=1)
# no reward for zero command
rew_airTime *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
return rew_airTime
def _reward_tracking_lin_vel(self, data, commands: np.ndarray):
# Tracking of linear velocity commands (xy axes)
lin_vel_error = np.sum(np.square(commands[:, :2] - self.get_local_linvel(data)[:, :2]), axis=1)
return np.exp(-lin_vel_error / self.cfg.reward_config.tracking_sigma)
def _reward_tracking_ang_vel(self, data, commands: np.ndarray):
# Tracking of angular velocity commands (yaw)
ang_vel_error = np.square(commands[:, 2] - self.get_gyro(data)[:, 2])
return np.exp(-ang_vel_error / self.cfg.reward_config.tracking_sigma)
def _reward_stand_still(self, data, commands: np.ndarray):
# Penalize motion at zero commands
return np.sum(np.abs(self.get_dof_pos(data) - self.default_angles), axis=1) * (
np.linalg.norm(commands, axis=1) < 0.1
)
def _reward_hip_pos(self, data, commands: np.ndarray):
return (0.8 - np.abs(commands[:, 1])) * np.sum(
np.square(self.get_dof_pos(data)[:, self.hip_indices] - self.default_angles[self.hip_indices]),
axis=1,
)
def _reward_calf_pos(self, data, commands: np.ndarray):
return (0.8 - np.abs(commands[:, 1])) * np.sum(
np.square(self.get_dof_pos(data)[:, self.calf_indices] - self.default_angles[self.calf_indices]),
axis=1,
)

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/calf.stl Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:ad3685e132454fd64e4ca3e5089f0d812fb53958e95c14913a22e3193e9e9ed8
size 1447984

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/hip.stl Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:b6599b9dd2751ebd87449203223f5683a6c235268b76fcc443e00e0d974a06c0
size 1505284

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/thigh.stl Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:bbf7c0a441323088c59276dbec8a32e708a1c61f42d30a6f6775176366ea23f8
size 793584

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/thigh_mirror.stl Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:1a90b89de47a57cb12a2d703f6b008a6ace8a1890a381c27723296fd0d00b9d8
size 950834

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/trunk.stl Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:051a57995dc6f892fc9bc53cd94e9a894f3dc209fd646e94d55d764150013970
size 5619284

238
motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml Normal file
View File

@@ -0,0 +1,238 @@
<mujoco model="go1">
<compiler angle="radian" autolimits="true"/>
<option timestep="0.005" integrator="Euler" iterations="60">
<flag eulerdamp="disable"/>
</option>
<custom>
<numeric data="30" name="max_contact_points"/>
<numeric data="12" name="max_geom_pairs"/>
</custom>
<default>
<default class="go1">
<geom condim="1"/>
<joint axis="0 1 0" armature="0.005" damping="0.5"/>
<!-- <position forcerange="-23.7 23.7" inheritrange="1" kp="35"/> -->
<default class="abduction">
<joint axis="1 0 0" range="-0.863 0.863" frictionloss="0.3"/>
</default>
<default class="hip">
<joint range="-0.686 4.501" frictionloss="0.3"/>
</default>
<default class="knee">
<joint range="-2.818 -0.888" frictionloss="1.0"/>
<position forcerange="-35.55 35.55"/>
</default>
<default class="visual">
<geom type="mesh" contype="0" conaffinity="0" group="2" material="dark"/>
</default>
<default class="collision">
<geom group="3" type="capsule"/>
<default class="hip_left1">
<geom size="0.046 0.02" pos="0 0.045 0" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="hip_left2">
<geom size="0.031 0.02" pos="0 0.065 0" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="hip_left3">
<geom size="0.046 0.02" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="hip_right1">
<geom size="0.046 0.02" pos="0 -0.045 0" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="hip_right2">
<geom size="0.031 0.02" pos="0 -0.065 0" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="hip_right3">
<geom size="0.046 0.02" quat="1 1 0 0" type="cylinder"/>
</default>
<default class="thigh1">
<geom size="0.015" fromto="-0.02 0 0 -0.02 0 -0.16"/>
</default>
<default class="thigh2">
<geom size="0.015" fromto="0 0 0 -0.02 0 -0.1"/>
</default>
<default class="thigh3">
<geom size="0.015" fromto="-0.02 0 -0.16 0 0 -0.2"/>
</default>
<default class="calf1">
<geom size="0.01" fromto="0 0 0 0.02 0 -0.13"/>
</default>
<default class="calf2">
<geom size="0.01" fromto="0.02 0 -0.13 0 0 -0.2"/>
</default>
<default class="foot">
<geom type="sphere" size="0.023" pos="0 0 -0.213" solimp="0.9 .99 0.001" priority="10" condim="3"/>
</default>
</default>
</default>
</default>
<asset>
<material name="dark" rgba="0.2 0.2 0.2 1"/>
<mesh class="go1" file="./assets/trunk.stl"/>
<mesh class="go1" file="./assets/hip.stl"/>
<mesh class="go1" file="./assets/thigh_mirror.stl"/>
<mesh class="go1" file="./assets/calf.stl"/>
<mesh class="go1" file="./assets/thigh.stl"/>
</asset>
<worldbody>
<body name="trunk" pos="0 0 0.4" childclass="go1">
<camera name="track" pos="0.846 -1.3 0.316" xyaxes="0.866 0.500 0.000 -0.171 0.296 0.940" mode="trackcom"/>
<camera name="top" pos="-1 0 1" xyaxes="0 -1 0 0.7 0 0.7" mode="trackcom"/>
<camera name="side" pos="0 -2 1" xyaxes="1 0 0 0 1 2" mode="trackcom"/>
<!-- <camera name="back" pos="-2.4 0 0.8" target="trunk" mode="targetbody" trackposspeed="2" trackrotspeed="10"/> -->
<site name="head" pos="0.3 0 0" rgba="1 0 0 1" size="0.02" group="5"/>
<inertial pos="0.0223 0.002 -0.0005" quat="-0.00342088 0.705204 0.000106698 0.708996" mass="5.204" diaginertia="0.0716565 0.0630105 0.0168101"/>
<freejoint/>
<geom class="visual" mesh="trunk"/>
<geom name='trunk' class="collision" quat="1 0 0 0" pos="0.24 0 0" size="0.05 0.05 0.05" type="box"/>
<!-- <geom name='trunk1' class="collision" quat="1 0 1 0" pos="0 -0.04 0" size="0.058 0.125" type="cylinder"/> -->
<!-- <geom name='trunk2' class="collision" quat="1 0 1 0" pos="0 +0.04 0" size="0.058 0.125" type="cylinder"/> -->
<site name="imu" pos="-0.01592 -0.06659 -0.00617" group="5"/>
<body name="FR_hip" pos="0.1881 -0.04675 0">
<inertial pos="-0.0049166 0.00762615 -8.865e-05" quat="0.507341 0.514169 0.495027 0.482891" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
<joint class="abduction" name="FR_hip_joint"/>
<geom class="visual" mesh="hip" quat="1 0 0 0"/>
<!-- <geom name="fr_hip" class="hip_right1"/> -->
<body name="FR_thigh" pos="0 -0.08 0">
<inertial pos="-0.00304722 0.019315 -0.0305004" quat="0.65243 -0.0272313 0.0775126 0.753383" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
<joint class="hip" name="FR_thigh_joint"/>
<geom class="visual" mesh="thigh_mirror"/>
<!-- <geom name="fr_thigh1" class="thigh1"/> -->
<!-- <geom name="fr_thigh2" class="thigh2"/> -->
<!-- <geom name="fr_thigh3" class="thigh3"/> -->
<body name="FR_calf" pos="0 0 -0.213">
<inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
<joint class="knee" name="FR_calf_joint"/>
<geom class="visual" mesh="calf"/>
<!-- <geom name="fr_calf1" class="calf1"/> -->
<!-- <geom name="fr_calf2" class="calf2"/> -->
<geom name="FR_foot" class="foot"/>
<site name="FR" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
</body>
</body>
</body>
<body name="FL_hip" pos="0.1881 0.04675 0">
<inertial pos="-0.0049166 -0.00762615 -8.865e-05" quat="0.482891 0.495027 0.514169 0.507341" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
<joint class="abduction" name="FL_hip_joint"/>
<geom class="visual" mesh="hip"/>
<!-- <geom name="fl_hip" class="hip_left1"/> -->
<body name="FL_thigh" pos="0 0.08 0">
<inertial pos="-0.00304722 -0.019315 -0.0305004" quat="0.753383 0.0775126 -0.0272313 0.65243" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
<joint class="hip" name="FL_thigh_joint"/>
<geom class="visual" mesh="thigh"/>
<!-- <geom name="fl_thigh1" class="thigh1"/> -->
<!-- <geom name="fl_thigh2" class="thigh2"/> -->
<!-- <geom name="fl_thigh3" class="thigh3"/> -->
<body name="FL_calf" pos="0 0 -0.213">
<inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
<joint class="knee" name="FL_calf_joint"/>
<geom class="visual" mesh="calf"/>
<!-- <geom name="fl_calf1" class="calf1"/> -->
<!-- <geom name="fl_calf2" class="calf2"/> -->
<geom name="FL_foot" class="foot"/>
<site name="FL" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
</body>
</body>
</body>
<body name="RR_hip" pos="-0.1881 -0.04675 0">
<inertial pos="0.0049166 0.00762615 -8.865e-05" quat="0.495027 0.482891 0.507341 0.514169" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
<joint class="abduction" name="RR_hip_joint"/>
<geom class="visual" quat="0 0 0 -1" mesh="hip"/>
<!-- <geom name="rr_hip" class="hip_right1"/> -->
<body name="RR_thigh" pos="0 -0.08 0">
<inertial pos="-0.00304722 0.019315 -0.0305004" quat="0.65243 -0.0272313 0.0775126 0.753383" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
<joint class="hip" name="RR_thigh_joint"/>
<geom class="visual" mesh="thigh_mirror"/>
<!-- <geom name="rr_thigh1" class="thigh1"/>
<geom name="rr_thigh2" class="thigh2"/>
<geom name="rr_thigh3" class="thigh3"/> -->
<body name="RR_calf" pos="0 0 -0.213">
<inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
<joint class="knee" name="RR_calf_joint"/>
<geom class="visual" mesh="calf"/>
<!-- <geom name="rr_calf1" class="calf1"/>
<geom name="rr_calf2" class="calf2"/> -->
<geom name="RR_foot" class="foot"/>
<site name="RR" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
</body>
</body>
</body>
<body name="RL_hip" pos="-0.1881 0.04675 0">
<inertial pos="0.0049166 -0.00762615 -8.865e-05" quat="0.514169 0.507341 0.482891 0.495027" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
<joint class="abduction" name="RL_hip_joint"/>
<geom class="visual" quat="0 0 1 0" mesh="hip"/>
<!-- <geom name="rl_hip" class="hip_left1"/> -->
<body name="RL_thigh" pos="0 0.08 0">
<inertial pos="-0.00304722 -0.019315 -0.0305004" quat="0.753383 0.0775126 -0.0272313 0.65243" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
<joint class="hip" name="RL_thigh_joint"/>
<geom class="visual" mesh="thigh"/>
<!-- <geom name="rl_thigh1" class="thigh1"/>
<geom name="rl_thigh2" class="thigh2"/>
<geom name="rl_thigh3" class="thigh3"/> -->
<body name="RL_calf" pos="0 0 -0.213">
<inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
<joint class="knee" name="RL_calf_joint"/>
<geom class="visual" mesh="calf"/>
<!-- <geom name="rl_calf1" class="calf1"/>
<geom name="rl_calf2" class="calf2"/> -->
<geom name="RL_foot" class="foot"/>
<site name="RL" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="FR_calf" joint="FR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="FL_calf" joint="FL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="RR_calf" joint="RR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="RL_calf" joint="RL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<!-- <position class="abduction" name="FR_hip" joint="FR_hip_joint" />
<position class="hip" name="FR_thigh" joint="FR_thigh_joint" />
<position class="knee" name="FR_calf" joint="FR_calf_joint" />
<position class="abduction" name="FL_hip" joint="FL_hip_joint" />
<position class="hip" name="FL_thigh" joint="FL_thigh_joint" />
<position class="knee" name="FL_calf" joint="FL_calf_joint" />
<position class="abduction" name="RR_hip" joint="RR_hip_joint" />
<position class="hip" name="RR_thigh" joint="RR_thigh_joint" />
<position class="knee" name="RR_calf" joint="RR_calf_joint" />
<position class="abduction" name="RL_hip" joint="RL_hip_joint" />
<position class="hip" name="RL_thigh" joint="RL_thigh_joint" />
<position class="knee" name="RL_calf" joint="RL_calf_joint" /> -->
</actuator>
<sensor>
<gyro site="imu" name="gyro"/>
<velocimeter site="imu" name="local_linvel"/>
<framepos objtype="site" objname="imu" name="position"/>
<!-- <framezaxis objtype="site" objname="imu" name="upvector"/>
<framexaxis objtype="site" objname="imu" name="forwardvector"/> -->
<framelinvel objtype="site" objname="imu" name="global_linvel"/>
<frameangvel objtype="site" objname="imu" name="global_angvel"/>
<framequat objtype="site" objname="imu" name="orientation"/>
<framelinvel objtype="site" objname="FR" name="FR_global_linvel"/>
<framelinvel objtype="site" objname="FL" name="FL_global_linvel"/>
<framelinvel objtype="site" objname="RR" name="RR_global_linvel"/>
<framelinvel objtype="site" objname="RL" name="RL_global_linvel"/>
<framepos objtype="site" objname="FR" name="FR_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="FL" name="FL_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="RR" name="RR_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="RL" name="RL_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="head" name="head_pos"/>
</sensor>
</mujoco>

58
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_motor_actuator.xml Normal file
View File

@@ -0,0 +1,58 @@
<mujoco model="go1 feetonly flat terrain scene">
<include file="go1_motor_actuator.xml"/>
<statistic center="0 0 0.1" extent="0.8" meansize="0.04"/>
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global azimuth="120" elevation="-20"/>
<map force="0.01"/>
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2"/>
<quality shadowsize="8192"/>
</visual>
<asset>
<texture type="skybox" builtin="gradient" rgb1="0.4314 0.5294 0.6431" rgb2="0 0 0" width="512" height="512"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.4314 0.5294 0.6431" rgb2="0.8157 0.8549 0.9059"
markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="1 1" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.01" type="plane" material="groundplane" contype="1" conaffinity="0" priority="1"
friction="0.6" condim="3"/>
</worldbody>
<keyframe>
<key name="home" qpos="
0 0 0.278
1 0 0 0
0.1 0.9 -1.8
-0.1 0.9 -1.8
0.1 0.9 -1.8
-0.1 0.9 -1.8"
ctrl="0.1 0.9 -1.8 -0.1 0.9 -1.8 0.1 0.9 -1.8 -0.1 0.9 -1.8"/>
<key name="home_higher" qpos="0 0 0.31 1 0 0 0 0 0.82 -1.63 0 0.82 -1.63 0 0.82 -1.63 0 0.82 -1.63"
ctrl="0 0.82 -1.63 0 0.82 -1.63 0 0.82 -1.63 0 0.82 -1.63"/>
<key
name="footstand"
qpos='
0 0 0.54
0.8 0 -0.8 0
0 0.82 -1.6 0 0.82 -1.68 0 1.82 -1.16 0.0 1.82 -1.16'
ctrl='0 0.82 -1.6 0 0.82 -1.68 0 1.82 -1.16 0.0 1.82 -1.16'/>
<key name="handstand"
qpos="0 0 0.54
0.8 0 0.8 0
0 -0.686 -1.16
0 -0.686 -1.16
0 1.7 -1.853
0 1.7 -1.853"
ctrl="0 -0.686 -1.16 0 -0.686 -1.16 0 1.7 -1.853 0 1.7 -1.853"/>
<key name="pre_recovery"
qpos="-0.0318481 -0.000215369 0.0579031 1 -2.70738e-05 6.06169e-05 0.000231261 -0.352275 1.18554 -2.80738 0.360892 1.1806 -2.80281 -0.381197 1.16812 -2.79123 0.391054 1.1622 -2.78576"
ctrl="-0.352275 1.18554 -2.80738 0.360892 1.1806 -2.80281 -0.381197 1.16812 -2.79123 0.391054 1.1622 -2.78576"/>
</keyframe>
</mujoco>

199
motrix_envs/src/motrix_envs/np/env.py Normal file
View File

@@ -0,0 +1,199 @@
# 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 abc
import dataclasses
from dataclasses import dataclass
import motrixsim as mtx
import numpy as np
from motrix_envs.base import ABEnv, EnvCfg
@dataclass
class NpEnvState:
data: mtx.SceneData
obs: np.ndarray
reward: np.ndarray
terminated: np.ndarray
truncated: np.ndarray
info: dict
@property
def done(self) -> np.ndarray:
"""
Check if the environment is done.
"""
return np.logical_or(self.terminated, self.truncated)
def replace(self, **updates) -> "NpEnvState":
return dataclasses.replace(self, **updates)
def validate(self):
num_envs = self.data.shape[0]
assert self.reward.shape == (num_envs,), self.reward.shape
assert self.terminated.shape == (num_envs,), self.terminated.shape
assert self.truncated.shape == (num_envs,), self.truncated.shape
class NpEnv(ABEnv):
_model: mtx.SceneModel
_cfg: EnvCfg
_state: NpEnvState = None
def __init__(self, cfg: EnvCfg, num_envs: int = 1):
self._cfg = cfg
self._num_envs = num_envs
self._model = mtx.load_model(cfg.model_file)
self._model.options.timestep = cfg.sim_dt
@property
def model(self) -> mtx.SceneModel:
"""
Get the scene model
"""
return self._model
@property
def state(self) -> NpEnvState:
"""
Get the current environment state
"""
return self._state
@property
def cfg(self) -> EnvCfg:
"""
Get the environment configuration
"""
return self._cfg
@property
def num_envs(self) -> int:
return self._num_envs
def init_state(self) -> NpEnvState:
"""
Create a new environment state
"""
obs = np.zeros((self._num_envs, self.observation_space.shape[0]), dtype=np.float32)
reward = np.zeros((self._num_envs,), dtype=np.float32)
terminated = np.ones((self._num_envs,), dtype=bool)
truncated = np.zeros((self._num_envs,), dtype=bool)
info = {"steps": np.zeros((self._num_envs,), dtype=np.uint64)}
data = mtx.SceneData(self._model, batch=[self._num_envs])
self._state = NpEnvState(data, obs, reward, terminated, truncated, info)
self._reset_done_envs()
self._state.validate()
return self._state
def _reset_done_envs(self):
"""
Reset the environments that are done
"""
state = self._state
done = state.done
assert done.shape == (self._num_envs,)
if not np.any(done):
return
np.putmask(state.info["steps"], done, 0)
data = state.data[done]
obs, info1 = self.reset(data)
state.obs[done] = obs
if info1:
def replace_dict_values(dst, new_values, mask):
for key, value in new_values.items():
if key not in dst:
dst[key] = value
else:
if isinstance(value, np.ndarray):
dst[key][mask] = value
elif isinstance(value, dict):
assert isinstance(dst[key], dict)
replace_dict_values(dst[key], value, mask)
replace_dict_values(state.info, info1, done)
def _update_truncate(self):
"""
Truncate the environments that have reached max episode length
"""
if not self._cfg.max_episode_steps:
return
self._state.truncated = self._state.info["steps"] >= self._cfg.max_episode_steps
@abc.abstractmethod
def apply_action(self, actions: np.ndarray, state: NpEnvState) -> NpEnvState:
"""
Apply the action to the environment
Args:
actions (np.ndarray): The actions to apply
state (NpEnvState): The environment state to apply the actions.
"""
@abc.abstractmethod
def update_state(self, state: NpEnvState) -> NpEnvState:
"""
Update the environment state after physics step
Args:
state (NpEnvState): The environment state to update
"""
@abc.abstractmethod
def reset(
self,
data: mtx.SceneData,
done: np.ndarray = None,
) -> tuple[np.ndarray, dict]:
"""
Reset the environment for the done envs
Args:
data (mtx.SceneData): The scene data to reset
done (Optional[np.ndarray]): A boolean array indicating which envs to reset. If None, reset all envs.
Returns:
tuple[np.ndarray, dict]: The initial observations and info after reset
"""
pass
def physics_step(self):
for _ in range(self._cfg.sim_substeps):
self._model.step(self._state.data)
def _prev_physics_step(self):
state = self._state
state.reward.fill(0.0)
state.terminated.fill(False)
state.truncated.fill(False)
def step(self, actions: np.ndarray) -> NpEnvState:
if self._state is None:
self.init_state()
self._prev_physics_step()
self._state = self.apply_action(actions, self._state)
assert self._state is not None, "apply_action must return a valid NpEnvState"
self.physics_step()
self._state = self.update_state(self._state)
self._state.info["steps"] += 1
self._update_truncate()
self._reset_done_envs()
return self._state

64
motrix_envs/src/motrix_envs/np/renderer.py Normal file
View File

@@ -0,0 +1,64 @@
# 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 numpy as np
from motrixsim.render import RenderApp, RenderSettings
from motrix_envs.np.env import NpEnv
class NpRenderer:
"""
The renderer for Np sim environments.
"""
_env: NpEnv
def __init__(self, env: NpEnv):
num_envs = env.num_envs
num_envs = 1 if num_envs is None else num_envs
spacing = 1.0
cols = int(np.ceil(np.sqrt(num_envs)))
offsets = []
for i in range(num_envs):
row = i // cols
col = i % cols
x = col * spacing
y = row * spacing
z = 0.0
offsets.append([x, y, z])
self._env = env
self._render = RenderApp()
settings = RenderSettings.performance()
settings.enable_shadow = False # disable shadow for better performance
self._render.launch(
env.model,
batch=num_envs,
render_offset=offsets,
render_settings=settings,
)
self._sync_render_data = True
self._render.system_camera.active = self._sync_render_data
def render(self) -> None:
"""
render the env
"""
self._render.sync(data=self._env.state.data if self._sync_render_data else None)
if self._render.input.is_key_just_pressed("space"):
self._sync_render_data = not self._sync_render_data
self._render.system_camera.active = self._sync_render_data

83
motrix_envs/src/motrix_envs/np/reward.py Normal file
View File

@@ -0,0 +1,83 @@
# 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 numpy as np
_DEFAULT_VALUE_AT_MARGIN = 0.1
def _sigmoids(x, value_at_1, sigmoid):
if sigmoid in ("cosine", "linear", "quadratic"):
if not 0 <= value_at_1 < 1:
raise ValueError(f"`value_at_1` must be nonnegative and smaller than 1, got {value_at_1}.")
else:
if not 0 < value_at_1 < 1:
raise ValueError(f"`value_at_1` must be strictly between 0 and 1, got {value_at_1}.")
if sigmoid == "gaussian":
scale = np.sqrt(-2 * np.log(value_at_1))
return np.exp(-0.5 * (x * scale) ** 2)
elif sigmoid == "hyperbolic":
scale = np.arccosh(1 / value_at_1)
return 1 / np.cosh(x * scale)
elif sigmoid == "long_tail":
scale = np.sqrt(1 / value_at_1 - 1)
return 1 / ((x * scale) ** 2 + 1)
elif sigmoid == "reciprocal":
scale = 1 / value_at_1 - 1
return 1 / (abs(x) * scale + 1)
elif sigmoid == "linear":
scale = 1 - value_at_1
scaled_x = x * scale
return np.where(abs(scaled_x) < 1, 1 - scaled_x, 0.0)
elif sigmoid == "quadratic":
scale = np.sqrt(1 - value_at_1)
scaled_x = x * scale
return np.where(abs(scaled_x) < 1, 1 - scaled_x**2, 0.0)
elif sigmoid == "tanh_squared":
scale = np.arctanh(np.sqrt(1 - value_at_1))
return 1 - np.tanh(x * scale) ** 2
else:
raise ValueError(f"Unknown sigmoid type {sigmoid!r}.")
def tolerance(
x: np.ndarray,
bounds: tuple[float, float] = (0.0, 0.0),
margin: float = 0.0,
sigmoid: str = "gaussian",
value_at_margin: float = _DEFAULT_VALUE_AT_MARGIN,
) -> np.ndarray:
lower, upper = bounds
if lower > upper:
raise ValueError("lower bound must be less than upper bound")
if margin < 0:
raise ValueError("margin must be non-negative")
in_bounds = np.logical_and(lower <= x, x <= upper)
if margin == 0:
value = np.where(in_bounds, 1.0, 0.0)
else:
d = np.where(x < lower, lower - x, x - upper) / margin
value = np.where(in_bounds, 1.0, _sigmoids(d, value_at_margin, sigmoid))
return value

171
motrix_envs/src/motrix_envs/registry.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.
# ==============================================================================
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional, Type, TypeVar
from motrix_envs.base import ABEnv, EnvCfg
TEnvCfg = TypeVar("TEnvCfg", bound=EnvCfg)
@dataclass
class EnvMeta:
env_cfg_cls: Type[EnvCfg]
env_cls_dict: Dict[str, Type[ABEnv]] = field(default_factory=dict)
def available_sim_backend(self) -> Optional[str]:
"""Return the first available simulation backend."""
return next(iter(self.env_cls_dict), None)
def support_sim_backend(self, sim_backend: str) -> bool:
"""Check if the environment supports a specific simulation backend."""
return sim_backend in self.env_cls_dict
_envs: Dict[str, EnvMeta] = {}
def contains(name: str) -> bool:
"""Check if an environment configuration is registered."""
return name in _envs
def register_env_config(name: str, env_cfg_cls: Type[EnvCfg]):
"""Register an environment configuration class with a name."""
if name in _envs.keys():
raise ValueError(f"Environment '{name}' is already registered.")
_envs[name] = EnvMeta(env_cfg_cls=env_cfg_cls)
def envcfg(name: str) -> Callable[[Type[TEnvCfg]], Type[TEnvCfg]]:
"""
Decorator to register an environment configuration class with a name.
Usage:
@register_env_config_decorator("my-env")
@dataclass
class MyEnvCfg(EnvCfg):
...
"""
def decorator(cls: Type[TEnvCfg]) -> Type[TEnvCfg]:
register_env_config(name, cls)
return cls
return decorator
def register_env(name: str, env_cls: Type[ABEnv], sim_backend: str):
"""Register an environment class with a name and simulation backend."""
if sim_backend not in ["np"]:
raise ValueError(f"Unsupported simulation backend: {sim_backend}. Only 'np' is supported yet.")
if name not in _envs:
raise ValueError(f"Environment '{name}' is not registered. Please register the config first.")
if sim_backend in _envs[name].env_cls_dict:
raise ValueError(f"Environment '{name}' with sim backend '{sim_backend}' is already registered.")
_envs[name].env_cls_dict[sim_backend] = env_cls
def env(name: str, sim_backend: str) -> Callable[[Type[ABEnv]], Type[ABEnv]]:
"""
Decorator to register an environment class with a name and simulation backend.
Usage:
@register_env_decorator("my-env", "np")
class MyEnv(ABEnv):
...
"""
def decorator(cls: Type[ABEnv]) -> Type[ABEnv]:
register_env(name, cls, sim_backend)
return cls
return decorator
def find_available_sim_backend(env_name: str) -> str:
"""Find the first available simulation backend for an environment."""
if env_name not in _envs:
raise ValueError(f"Environment '{env_name}' is not registered.")
meta: EnvMeta = _envs[env_name]
backend = meta.available_sim_backend()
if backend is None:
raise ValueError(f"Environment '{env_name}' does not support any simulation backend.")
return backend
def make(
name: str,
sim_backend: Optional[str] = None,
env_cfg_override: Optional[Dict[str, Any]] = None,
num_envs: int = 1,
) -> ABEnv:
"""
Create an environment instance by name.
Args:
name: Environment name
sim_backend: Simulation backend ("np"). If None, uses first available.
env_cfg_override: Dictionary of config overrides
num_envs: Number of environments to create
Returns:
Environment instance
"""
if name not in _envs:
raise ValueError(f"Environment '{name}' is not registered.")
meta: EnvMeta = _envs[name]
# Create environment config
env_cfg = meta.env_cfg_cls()
if env_cfg_override is not None:
for key, value in env_cfg_override.items():
if hasattr(env_cfg, key):
setattr(env_cfg, key, value)
else:
raise ValueError(f"Config class '{env_cfg.__class__.__name__}' has no attribute '{key}'")
# Validate config
env_cfg.validate()
# Select simulation backend
if sim_backend is None:
sim_backend = meta.available_sim_backend()
if sim_backend is None:
raise ValueError(f"Environment '{name}' does not support any simulation backend.")
if not meta.support_sim_backend(sim_backend):
raise ValueError(f"Environment '{name}' does not support simulation backend '{sim_backend}'.")
# Create environment instance
env_cls = meta.env_cls_dict[sim_backend]
return env_cls(env_cfg, num_envs=num_envs)
def list_registered_envs() -> Dict[str, Dict[str, Any]]:
"""List all registered environments with their available backends."""
result = {}
for name, meta in _envs.items():
result[name] = {
"config_class": meta.env_cfg_cls.__name__,
"available_backends": list(meta.env_cls_dict.keys()),
}
return result