chore: release v0.2.0

This commit is contained in:
motphys-developers
2026-02-10 08:08:11 +00:00
parent dbfa9e31fa
commit b568ac5600
123 changed files with 9732 additions and 497 deletions

View File

@@ -4,12 +4,12 @@ build-backend = "uv_build"
[project]
name = "motrix-envs"
version = "0.1.0"
version = "0.2.0"
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.5.0b2",
"motrixsim>=0.6.0b1",
]

View File

@@ -13,4 +13,16 @@
# limitations under the License.
# ==============================================================================
from . import bounce_ball, cartpole, cheetah, hopper, reacher, walker # noqa: F401 import to register envs
from . import ( # noqa: F401 import to register envs
acrobot,
bounce_ball,
cartpole,
cheetah,
finger,
hopper,
humanoid,
manipulator,
pendulum,
reacher,
walker,
)

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

View File

@@ -0,0 +1,49 @@
<!--
Based on Coulomb's [1] rather than Spong's [2] model.
[1] Coulom, Rémi. Reinforcement learning using neural networks, with applications to motor control.
Diss. Institut National Polytechnique de Grenoble-INPG, 2002.
[2] Spong, Mark W. "The swing up control problem for the acrobot."
IEEE control systems 15, no. 1 (1995): 49-55.
-->
<mujoco model="acrobot">
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<default>
<joint damping="0.05"/>
<geom type="capsule" mass="1"/>
</default>
<option timestep="0.01" integrator="RK4" gravity="0 0 -9.81">
<flag contact="disable" energy="enable"/>
</option>
<asset>
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0"
width="512" height="512"/>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png"/>
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4"/>
</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="motphys-ground" pos="0 0 -1"/>
<site name="target" type="sphere" pos="0 0 4" size="0.2" material="target" group="3"/>
<camera name="fixed" pos="0 -6 2" zaxis="0 -1 0"/>
<camera name="lookat" mode="targetbodycom" target="upper_arm" pos="0 -2 3"/>
<body name="upper_arm" pos="0 0 2">
<joint name="shoulder" type="hinge" axis="0 1 0"/>
<geom name="upper_arm_decoration" material="decoration" type="cylinder" fromto="0 -.06 0 0 .06 0" size="0.051" mass="0"/>
<geom name="upper_arm" fromto="0 0 0 0 0 1" size="0.05" material="self"/>
<body name="lower_arm" pos="0 0 1">
<joint name="elbow" type="hinge" axis="0 1 0"/>
<geom name="lower_arm" fromto="0 0 0 0 0 1" size="0.049" material="self"/>
<site name="tip" pos="0 0 1" size="0.01"/>
</body>
</body>
</worldbody>
<actuator>
<motor name="elbow" joint="elbow" gear="30" ctrllimited="true" ctrlrange="-1 1"/>
</actuator>
</mujoco>

View File

@@ -0,0 +1,152 @@
# 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 import reward
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import AcrobotEnvCfg
@registry.env("acrobot", "np")
class AcrobotEnv(NpEnv):
_cfg: AcrobotEnvCfg
def __init__(self, cfg: AcrobotEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
self._action_space = gym.spaces.Box(-1.0, 1.0, (1,), dtype=np.float32)
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (6,), dtype=np.float32)
self._num_dof_pos = self._model.num_dof_pos
self._num_dof_vel = self._model.num_dof_vel
self._tip = self._model.get_site("tip")
self._target = self._model.get_site("target")
self._upper_arm = self._model.get_link("upper_arm")
self._lower_arm = self._model.get_link("lower_arm")
self._target_radius = 0.2
self._step_count = np.zeros(self._num_envs, dtype=np.int32)
self._max_steps = int(cfg.max_episode_seconds / cfg.ctrl_dt)
@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):
actions = np.clip(actions, -1.0, 1.0)
state.data.actuator_ctrls = actions
return state
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
dof_pos = data.dof_pos
shoulder_angle = dof_pos[:, 0]
elbow_angle = dof_pos[:, 1]
upper_arm_horizontal = np.cos(shoulder_angle)
upper_arm_vertical = np.sin(shoulder_angle)
total_angle = shoulder_angle + elbow_angle
lower_arm_horizontal = np.cos(total_angle)
lower_arm_vertical = np.sin(total_angle)
dof_vel = data.dof_vel
obs = np.concatenate(
[
upper_arm_horizontal.reshape(-1, 1),
lower_arm_horizontal.reshape(-1, 1),
upper_arm_vertical.reshape(-1, 1),
lower_arm_vertical.reshape(-1, 1),
dof_vel,
],
axis=-1,
)
return obs
def update_state(self, state: NpEnvState) -> NpEnvState:
data = state.data
obs = self._get_obs(data)
tip_pos = self._tip.get_pose(data)
target_pos = self._target.get_pose(data)
dist_to_target = np.linalg.norm(tip_pos[:, :3] - target_pos[:, :3], axis=-1)
base_rwd = reward.tolerance(
dist_to_target,
bounds=(0, self._target_radius),
margin=0,
value_at_margin=0.0,
sigmoid="linear",
)
in_target = dist_to_target < self._target_radius
continuous_reward = 0.1 * in_target
distance_reward = 0.3 * (1.0 - np.clip(dist_to_target / 2.0, 0, 1.0))
dof_vel = data.dof_vel
vel_magnitude = np.mean(np.abs(dof_vel), axis=-1)
velocity_penalty = 0.01 * np.maximum(0, vel_magnitude - 2.0)
rwd = base_rwd + continuous_reward + distance_reward - velocity_penalty
self._step_count += 1
terminated = np.zeros((self._num_envs,), dtype=bool)
terminated = np.logical_or(self._step_count >= self._max_steps, terminated)
terminated = np.logical_or(np.isnan(obs).any(axis=-1), terminated)
state.obs = obs
state.reward = rwd
state.terminated = terminated
return state
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
data.reset(self._model)
num_reset = data.shape[0]
shoulder_angle = np.random.uniform(-np.pi, np.pi, size=num_reset).astype(np.float32)
elbow_angle = np.random.uniform(-np.pi, np.pi, size=num_reset).astype(np.float32)
dof_pos = np.stack([shoulder_angle, elbow_angle], axis=-1)
dof_vel = np.zeros((*data.shape, self._num_dof_vel), dtype=np.float32)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
obs = self._get_obs(data)
return obs, {}
def _reset_done_envs(self):
"""
Reset the environments that are done
"""
super()._reset_done_envs()
done = self._state.done
if np.any(done):
self._step_count[done] = 0

View File

@@ -0,0 +1,33 @@
# 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__) + "/acrobot.xml"
@registry.envcfg("acrobot")
@dataclass
class AcrobotEnvCfg(EnvCfg):
model_file: str = model_file
reset_noise_scale: float = 0.1
max_episode_seconds: float = 10.0
render_spacing: float = 2.0
sim_dt: float = 0.01
ctrl_dt: float = 0.02

View File

@@ -31,8 +31,34 @@
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom type="plane" size="0 0 .01" material="motphys-ground"/>
<!-- Target height visualization using mocap body (can be moved at runtime) -->
<body name="target_height_marker" mocap="true" pos="0.58856 0 0.5">
<geom name="target_height_visual"
type="cylinder"
size="0.1 0.002"
rgba="0 1 0 0.5"
contype="0"
conaffinity="0"/>
<geom name="target_center_dot"
type="cylinder"
size="0.01 0.002"
rgba="0 0.6 0 0.9"
contype="0"
conaffinity="0"/>
</body>
<!-- Paddle home position marker -->
<body name="paddle_home_marker" mocap="true" pos="0.5857 -0.0082 0.2">
<geom name="paddle_home_visual"
type="box"
size="0.1 0.1 0.003"
rgba="1 1 0.6 0.4"
contype="0"
conaffinity="0"/>
</body>
<geom type="mesh" rgba="0.25098 0.25098 0.25098 1" mesh="base_link" class="visual"/>
<body name="Link1" pos="0 0 0.1">
<body name="Link1" pos="0 0 0.33">
<inertial pos="0.00022014 -7.0626e-06 -0.10379" quat="0.998051 -0.0622621 0.00360518 0.00215072" mass="0.98482" diaginertia="0.00658905 0.00502396 0.00428069"/>
<joint name="Joint1" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" actuatorfrcrange="-300 300"/>
<geom type="mesh" rgba="1 0.69804 0 1" mesh="Link1" class="visual"/>
@@ -68,7 +94,7 @@
<body name="ball_link" pos="0 0 0.07">
<freejoint/>
<inertial pos="0 0 0" mass="0.0027" diaginertia="1 1 1"/>
<geom size="0.019" contype="1" rgba="0 1 0 1" solref="1 0"/>
<geom size="0.019" contype="1" rgba="1 0 0 1" solref="1 0"/>
</body>
</worldbody>

View File

@@ -30,36 +30,71 @@ class BounceBallEnv(NpEnv):
def __init__(self, cfg: BounceBallEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
# Action space: 6D normalized paddle velocity (dx, dy, dz, dr_x, dr_y, dr_z)
# Action space: 6D joint position control
self._action_space = gym.spaces.Box(-1.0, 1.0, (6,), dtype=np.float32)
# Observation space: simplified version using only DOF information
# DOF pos (13) + DOF vel (12) = 25 (this includes ball state implicitly)
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (25,), dtype=np.float32)
# Observation space: joint states + paddle position + target height (29D)
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (29,), dtype=np.float32)
self._num_dof_pos = self._model.num_dof_pos
self._num_dof_vel = self._model.num_dof_vel
# Initial arm joint positions (degrees converted to radians)
# Initial arm joint positions
self._init_arm_qpos = np.array(self._cfg.arm_init_qpos, dtype=np.float32) * np.pi / 180.0
self._init_dof_vel = np.zeros(self._model.num_dof_vel, dtype=np.float32)
# Initialize full DOF positions (6 arm joints + 7 for ball free joint)
# Full DOF positions (6 arm joints + 7 ball free joint)
self._init_dof_pos = np.zeros(self._model.num_dof_pos, dtype=np.float32)
self._init_dof_pos[:6] = self._init_arm_qpos
# Get body and geom IDs
self._paddle_geom_id = self._model.geom_names.index("blocker")
# Body and geom references
self._paddle_geom = self._model.get_geom("blocker")
self._ball_body_id = self._model.body_names.index("ball_link")
# Action scaling parameters
# Mocap bodies for visual markers
self._target_marker_body = self._model.get_body("target_height_marker")
assert self._target_marker_body.is_mocap, "target_height_marker must be a mocap body"
self._paddle_home_marker_body = self._model.get_body("paddle_home_marker")
assert self._paddle_home_marker_body.is_mocap, "paddle_home_marker must be a mocap body"
# Action scaling
self._action_scale = np.array(self._cfg.action_scale, dtype=np.float32)
self._action_bias = np.array(self._cfg.action_bias, dtype=np.float32)
# Track ball initial position for reset
# Ball initial conditions
self._ball_init_pos = np.array(self._cfg.ball_init_pos, dtype=np.float32)
self._ball_init_vel = np.array(self._cfg.ball_init_vel, dtype=np.float32)
# Constants for marker poses
self._ball_radius = 0.019 # Ball radius in meters
# Target marker base pose: [x, y, z_placeholder, qx, qy, qz, qw]
self._target_marker_base_pose = np.array(
[
self._cfg.target_ball_x,
self._cfg.target_ball_y,
0.0, # z will be set per environment
0.0,
0.0,
0.0,
1.0, # identity quaternion
],
dtype=np.float32,
)
# Paddle home marker pose: [x, y, z, qx, qy, qz, qw]
self._paddle_home_marker_pose = np.array(
[
self._cfg.target_ball_x,
self._cfg.target_ball_y,
self._cfg.paddle_home_position_z,
0.0,
0.0,
0.0,
1.0, # identity quaternion
],
dtype=np.float32,
)
@property
def observation_space(self):
return self._observation_space
@@ -69,118 +104,361 @@ class BounceBallEnv(NpEnv):
return self._action_space
def _denormalize_action(self, action: np.ndarray) -> np.ndarray:
"""Denormalize action to get actual paddle velocity changes"""
"""Denormalize action from [-1, 1] to joint position changes"""
return self._action_scale * action + self._action_bias
def _compute_observation(self, data: mtx.SceneData) -> np.ndarray:
"""Compute 25-dimensional observation vector from DOF states"""
# Use DOF positions and velocities directly
def _compute_observation(self, data: mtx.SceneData, target_heights: np.ndarray) -> np.ndarray:
"""Compute observation: joint states + paddle position + target height (29D)"""
dof_pos = data.dof_pos
dof_vel = data.dof_vel
# Concatenate DOF positions (13) and velocities (12)
obs = np.concatenate([dof_pos, dof_vel], axis=-1)
# Get paddle position
paddle_pose = self._paddle_geom.get_pose(data)
paddle_xyz = paddle_pose[:, :3]
# Concatenate: DOF pos (13) + DOF vel (12) + paddle xyz (3) + target height (1)
obs = np.concatenate([dof_pos, dof_vel, paddle_xyz, target_heights[:, np.newaxis]], axis=-1)
return obs.astype(np.float32)
def _compute_reward(
self, obs: np.ndarray, data: mtx.SceneData = None, consecutive_bounces: np.ndarray = None
) -> np.ndarray:
"""Compute reward based on ball height, position, and controlled upward velocity"""
# Extract ball position and velocity from DOF
ball_x = obs[:, 6] # Ball x position
ball_z = obs[:, 8] # Ball z position
self,
obs: np.ndarray,
data: mtx.SceneData = None,
consecutive_bounces: np.ndarray = None,
bounce_detected: np.ndarray = None,
target_heights: np.ndarray = None,
current_actions: np.ndarray = None,
last_actions: np.ndarray = None,
) -> tuple:
"""
Compute reward based on ball position, velocity, and paddle alignment.
ball_vz = obs[:, 13 + 8] # Ball z velocity (13 pos + 8 vel)
The reward function uses a composite design with multiple reward and penalty terms
to guide the robot to learn a stable ball bouncing strategy.
Returns:
tuple: (total_reward, reward_details) where reward_details contains
individual reward components for analysis.
"""
# Extract ball state
ball_x = obs[:, 6]
ball_y = obs[:, 7]
ball_z = obs[:, 8]
ball_vz = obs[:, 13 + 8]
# Extract paddle position
paddle_xy = obs[:, 25:27]
paddle_z = obs[:, 27]
# Target positions
target_ball_x = 0.58856 # Target x position
target_height = self._cfg.target_ball_height
tolerance = self._cfg.height_tolerance
target_ball_x = self._cfg.target_ball_x
target_ball_y = self._cfg.target_ball_y
target_height = target_heights
# 1. Position control reward - MOST IMPORTANT for keeping ball centered
# Strong reward for ball being at the right x position (paddle center)
# Physics constant
g = self._cfg.gravity
# ============================================================================
# 1. Horizontal position reward (weighted by vertical distance)
# Core reward ensuring ball stays directly above paddle
# ============================================================================
x_position_error = np.abs(ball_x - target_ball_x)
x_position_reward = np.exp(-(x_position_error**2) / (2 * 0.05**2)) # Tight tolerance for x position
y_position_error = np.abs(ball_y - target_ball_y)
xy_position_error = np.sqrt(x_position_error**2 + y_position_error**2)
# 2. Height-based reward - less important than position control
vertical_dist = np.abs(ball_z - paddle_z)
vertical_weight = np.exp(-vertical_dist / self._cfg.vertical_weight_scale)
weighted_horizontal_scale = self._cfg.weighted_horizontal_base_scale * (
1.0 + self._cfg.weighted_horizontal_weight_factor * vertical_weight
)
weighted_position_reward = np.exp(-(xy_position_error**2) / (2 * weighted_horizontal_scale**2))
# ============================================================================
# 2. Out-of-position penalty
# Strong penalty for severe deviation to prevent ball flying out of control
# ============================================================================
out_of_position_penalty = -2.0 / (
1.0
+ np.exp(-(xy_position_error - self._cfg.out_of_position_threshold) / self._cfg.out_of_position_sharpness)
)
# ============================================================================
# 3. Velocity matching reward
# Based on projectile motion physics, encourages ball trajectory to have
# desired velocity at target height for stable control
# ============================================================================
height_diff = target_height - ball_z
desired_velocity_at_target = self._cfg.desired_velocity_at_target
# Calculate velocity at target height using physics
upward_below_condition = (ball_vz > 0) & (ball_z < target_height)
velocity_squared_at_target_upward = np.where(upward_below_condition, ball_vz**2 - 2 * g * height_diff, 0.0)
velocity_at_target_upward = np.sqrt(np.maximum(0, velocity_squared_at_target_upward))
downward_above_condition = (ball_vz < 0) & (ball_z > target_height)
velocity_squared_at_target_downward = np.where(
downward_above_condition, ball_vz**2 + 2 * g * np.abs(height_diff), 0.0
)
velocity_at_target_downward = -np.sqrt(np.maximum(0, velocity_squared_at_target_downward))
near_target_condition = np.abs(height_diff) < 0.05
velocity_at_target_near = np.where(near_target_condition, ball_vz, 0.0)
# Smooth combination
upward_motion = 1.0 / (1.0 + np.exp(-ball_vz / 0.2))
below_target = 1.0 / (1.0 + np.exp(-height_diff / 0.02))
downward_motion = 1.0 - upward_motion
above_target = 1.0 - below_target
at_target_weight = np.exp(-(height_diff**2) / (2 * 0.01**2))
velocity_at_target = (
velocity_at_target_upward * upward_motion * below_target
+ velocity_at_target_downward * downward_motion * above_target
+ velocity_at_target_near * at_target_weight
)
velocity_error = np.abs(velocity_at_target - desired_velocity_at_target)
velocity_matching_reward = np.exp(-(velocity_error**2) / (2 * self._cfg.velocity_error_sigma**2))
# ============================================================================
# 4. Height reward
# Directly encourages ball to approach target height, core task objective
# ============================================================================
height_error = np.abs(ball_z - target_height)
height_reward = np.exp(-(height_error**2) / (2 * tolerance**2))
height_reward = np.exp(-(height_error**2) / (2 * self._cfg.height_error_sigma**2))
# 3. Controlled upward velocity reward - only when ball is in good position
# Only reward upward velocity when ball is well-positioned horizontally
well_positioned = x_position_error < 0.02 # Ball must be very close to target x
controlled_upward_reward = np.where(
well_positioned & (ball_vz > 0.1) & (ball_vz < 1.5), # Reasonable upward velocity
np.clip(ball_vz * 1.5, 0.0, 1.5), # Reduced scale
0.0,
# Height progress bonus
height_progress_bonus = (
np.maximum(0, ball_z - self._cfg.height_progress_threshold) * self._cfg.height_progress_scale
)
# 4. Strong penalty for being out of position horizontally
out_of_position_penalty = np.where(
x_position_error > 0.1,
-2.0, # Heavy penalty for being far from center
0.0,
# ============================================================================
# 5. Controlled upward velocity reward
# Only rewards upward velocity when ball position is good, avoiding random hitting
# Guides strategy to learn precise hitting force
# ============================================================================
positioning_quality = np.exp(-(xy_position_error**2) / (2 * self._cfg.positioning_quality_sigma**2))
ideal_launch_velocity = np.sqrt(2 * g * np.maximum(0, height_diff))
ideal_launch_velocity = np.clip(
ideal_launch_velocity, self._cfg.ideal_velocity_min, self._cfg.ideal_velocity_max
)
# 5. Velocity penalties - discourage excessive speeds
excessive_upward_penalty = np.where(ball_vz > 2.0, -1.0, 0.0)
downward_velocity_penalty = np.where(ball_vz < -0.5, -np.clip(-ball_vz * 0.3, 0.0, 0.5), 0.0)
# 6. Position-based penalties (reduced)
overshoot_penalty = np.where(ball_z > target_height + tolerance, -0.3, 0.0)
undershoot_penalty = np.where(ball_z < 0.1, -0.5, 0.0)
# 7. Consecutive bounces reward - only when position is good
good_position_for_bounce = x_position_error < 0.05 # Reasonable position for bouncing
consecutive_bounces_reward = np.where(
good_position_for_bounce & (consecutive_bounces > 0),
np.log(consecutive_bounces + 1) * 0.3, # Reduced scale
0.0,
upward_velocity_quality = np.exp(
-((ball_vz - ideal_launch_velocity) ** 2) / (2 * self._cfg.upward_velocity_sigma**2)
)
# Bonus for high bounce counts (only when well-positioned)
high_bounce_bonus = np.where(
good_position_for_bounce & (consecutive_bounces >= 3),
consecutive_bounces * 0.1, # Reduced bonus
0.0,
upward_mask = 1.0 / (1.0 + np.exp(-ball_vz / self._cfg.upward_mask_scale))
controlled_upward_reward = (
positioning_quality
* upward_velocity_quality
* upward_mask
* np.clip(ball_vz * 1.0, 0.0, self._cfg.controlled_upward_clip_max)
)
# Combine all rewards with corrected priorities
# ============================================================================
# 6. Velocity penalties
# Prevents ball velocity from being too fast or falling freely
# Ensures ball motion stays within controllable range
# ============================================================================
excessive_upward_penalty = -1.0 / (
1.0 + np.exp(-(ball_vz - self._cfg.excessive_upward_threshold) / self._cfg.excessive_upward_sharpness)
)
downward_penalty_magnitude = -ball_vz * np.clip(
-ball_vz * self._cfg.downward_velocity_scale, 0.0, self._cfg.downward_velocity_clip_max
)
downward_penalty_trigger = 1.0 / (1.0 + np.exp((ball_vz - self._cfg.downward_velocity_threshold) / 0.2))
downward_velocity_penalty = downward_penalty_magnitude * downward_penalty_trigger
# ============================================================================
# 7. Consecutive bounces reward
# Encourages multiple consecutive successful bounces for stable long-term control
# Uses logarithmic function to avoid infinite reward growth
# ============================================================================
bounce_positioning_quality = np.exp(-(xy_position_error**2) / (2 * self._cfg.bounce_positioning_sigma**2))
bounce_log_reward = np.log(consecutive_bounces.astype(np.float32) + 1.0) * self._cfg.bounce_log_scale
bounce_activation = (consecutive_bounces > 0).astype(np.float32) * bounce_positioning_quality
consecutive_bounces_reward = bounce_log_reward * bounce_activation
# High bounce count bonus
high_bounce_activation = 1.0 / (
1.0
+ np.exp(
-(consecutive_bounces.astype(np.float32) - self._cfg.high_bounce_threshold)
/ self._cfg.high_bounce_sharpness
)
)
high_bounce_bonus = (
consecutive_bounces.astype(np.float32)
* self._cfg.high_bounce_scale
* bounce_positioning_quality
* high_bounce_activation
)
# ============================================================================
# 8. Paddle-ball horizontal alignment
# Encourages paddle to actively move directly below ball
# Extra reward at bounce moment to reinforce correct hitting behavior
# ============================================================================
if bounce_detected is None:
bounce_detected = np.zeros(ball_x.shape[0], dtype=bool)
ball_xy = np.stack([ball_x, ball_y], axis=1)
paddle_ball_xy_error = np.linalg.norm(ball_xy - paddle_xy, axis=1)
vertical_proximity_weight = np.exp(-vertical_dist / self._cfg.vertical_proximity_scale)
paddle_alignment_quality = np.exp(-(paddle_ball_xy_error**2) / (2 * self._cfg.paddle_alignment_sigma**2)) * (
1.0 + self._cfg.paddle_alignment_weight_factor * vertical_proximity_weight
)
bounce_boost = bounce_detected.astype(np.float32) * self._cfg.bounce_boost_factor + 1.0
paddle_center_reward = paddle_alignment_quality * bounce_boost * self._cfg.paddle_center_scale
# ============================================================================
# 9. Paddle home position reward
# Encourages paddle to return to home position when ball is far away
# Makes paddle motion more energy-efficient and natural
# ============================================================================
paddle_home_z = self._cfg.paddle_home_position_z
paddle_height_deviation = np.abs(paddle_z - paddle_home_z)
# Distance-based dynamic factor
distance_factor = 1.0 + self._cfg.distance_factor_scale / (
1.0 + np.exp(-(vertical_dist - self._cfg.distance_factor_threshold) / self._cfg.distance_factor_sharpness)
)
home_position_reward = (
np.exp(-(paddle_height_deviation**2) / (2 * self._cfg.home_position_sigma**2)) * distance_factor
)
# Height violation penalty
max_deviation = self._cfg.max_paddle_height_deviation
height_violation = np.maximum(0, paddle_height_deviation - max_deviation)
height_violation_penalty = -height_violation * self._cfg.height_violation_scale
# ============================================================================
# 10. Action and velocity penalties
# Penalizes drastic action changes and excessive joint velocities
# Encourages smooth and energy-efficient control
# ============================================================================
num_envs = obs.shape[0]
if current_actions is None:
current_actions = np.zeros((num_envs, 6), dtype=np.float32)
if last_actions is None:
last_actions = np.zeros((num_envs, 6), dtype=np.float32)
action_diff = current_actions - last_actions
action_penalty = np.sum(np.square(action_diff), axis=-1)
joint_vel = data.dof_vel[:, :6]
joint_vel_penalty = np.sum(np.square(joint_vel), axis=-1)
# ============================================================================
# Total reward
# ============================================================================
action_penalty_rate = self._cfg.action_penalty_rate
joint_vel_penalty_rate = self._cfg.joint_vel_penalty_rate
total_reward = (
x_position_reward * 2.0 # X position control (200%) - MOST IMPORTANT
+ controlled_upward_reward * 1.0 # Controlled upward velocity (100%)
+ height_reward * 0.3 # Height accuracy (30%) - less important
+ consecutive_bounces_reward * 1.0 # Consecutive bounces (100%) - reduced
+ high_bounce_bonus * 0.3 # High bounce bonus (30%) - reduced
+ out_of_position_penalty * 1.0 # Out of position penalty (100%)
+ excessive_upward_penalty * 1.0 # Excessive upward penalty (100%)
+ downward_velocity_penalty * 1.0 # Downward penalty (100%)
+ overshoot_penalty # Height overshoot penalty
+ undershoot_penalty # Height undershoot penalty
weighted_position_reward * self._cfg.weighted_position_weight
+ velocity_matching_reward * self._cfg.velocity_matching_weight
+ height_reward * self._cfg.height_reward_weight
+ height_progress_bonus * self._cfg.height_progress_weight
+ controlled_upward_reward * self._cfg.controlled_upward_weight
+ consecutive_bounces_reward * self._cfg.consecutive_bounces_weight
+ high_bounce_bonus * self._cfg.high_bounce_weight
+ paddle_center_reward * self._cfg.paddle_center_weight
+ home_position_reward * self._cfg.home_position_weight
+ out_of_position_penalty * self._cfg.out_of_position_weight
+ excessive_upward_penalty * self._cfg.excessive_upward_weight
+ downward_velocity_penalty * self._cfg.downward_velocity_weight
+ height_violation_penalty * self._cfg.height_violation_weight
- action_penalty_rate * action_penalty
- joint_vel_penalty_rate * joint_vel_penalty
)
return total_reward
# ============================================================================
# Reward details for debugging
# ============================================================================
reward_details = {
"x_position_error": x_position_error,
"y_position_error": y_position_error,
"xy_position_error": xy_position_error,
"ball_z": ball_z,
"paddle_z": paddle_z,
"vertical_dist": vertical_dist,
"vertical_weight": vertical_weight,
"ball_vz": ball_vz,
"height_diff": height_diff,
"height_error": height_error,
"velocity_at_target": velocity_at_target,
"desired_velocity_at_target": np.full_like(ball_vz, desired_velocity_at_target),
"velocity_error": velocity_error,
"ideal_launch_velocity": ideal_launch_velocity,
"positioning_quality": positioning_quality,
"upward_velocity_quality": upward_velocity_quality,
"bounce_positioning_quality": bounce_positioning_quality,
"paddle_ball_xy_error": paddle_ball_xy_error,
"vertical_proximity_weight": vertical_proximity_weight,
"bounce_detected": bounce_detected.astype(np.float32),
"consecutive_bounces": consecutive_bounces.astype(np.float32),
"action_penalty": action_penalty,
"joint_vel_penalty": joint_vel_penalty,
"distance_factor": distance_factor,
"weighted_position_reward": weighted_position_reward * 2.0,
"velocity_matching_reward": velocity_matching_reward * 2.0,
"height_reward": height_reward * 4.5,
"height_progress_bonus": height_progress_bonus * 1.0,
"controlled_upward_reward": controlled_upward_reward * 1.5,
"consecutive_bounces_reward": consecutive_bounces_reward * 0.8,
"high_bounce_bonus": high_bounce_bonus * 0.3,
"paddle_center_reward": paddle_center_reward * 0.6,
"home_position_reward": home_position_reward * 1.5,
"out_of_position_penalty": out_of_position_penalty * 1.0,
"excessive_upward_penalty": excessive_upward_penalty * 1.0,
"downward_velocity_penalty": downward_velocity_penalty * 1.0,
"height_violation_penalty": height_violation_penalty * 1.0,
"action_penalty_weighted": -action_penalty_rate * action_penalty,
"joint_vel_penalty_weighted": -joint_vel_penalty_rate * joint_vel_penalty,
"paddle_height_deviation": paddle_height_deviation,
"total_reward": total_reward,
}
def _compute_terminated(self, obs: np.ndarray) -> np.ndarray:
return total_reward, reward_details
def _compute_terminated(self, obs: np.ndarray, target_heights: np.ndarray) -> np.ndarray:
"""Check if episode should terminate based on DOF states"""
# Extract ball position from DOF (indices 6-8 for x,y,z)
ball_x = obs[:, 6] # Ball x position
ball_y = obs[:, 7] # Ball y position
ball_z = obs[:, 8] # Ball z position
# Target height from config
target_height = self._cfg.target_ball_height
# Terminate if ball falls below ground or goes significantly higher than target
terminated = (ball_z < 0.05) | (ball_z > target_height + 1.0)
terminated = (ball_z < 0.05) | (ball_z > target_heights + 1.0)
# Also terminate if ball goes too far horizontally
terminated |= np.abs(ball_x) > 1.5
terminated |= (np.abs(ball_x) > 1.5) | (np.abs(ball_y) > 1.5)
# Terminate if joint velocity is too high
# Limit: 360 degrees/second = 2*pi rad/s ≈ 6.28 rad/s
joint_vel = obs[:, 13:19] # Joint velocities (indices 13-18 for 6 arm joints)
max_joint_vel = 2.0 * np.pi # 360 degrees/second in radians
terminated |= np.abs(joint_vel).max(axis=-1) > max_joint_vel
return terminated
def apply_action(self, actions: np.ndarray, state: NpEnvState) -> NpEnvState:
"""Apply action to control paddle position"""
# Store last actions for penalty calculation
state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions))
state.info["current_actions"] = actions
# Get current joint positions
current_joint_pos = state.data.dof_pos[:, :6] # First 6 DOFs are arm joints
@@ -198,12 +476,15 @@ class BounceBallEnv(NpEnv):
"""Update state with new observations, rewards, and termination flags"""
data = state.data
# Compute observation
obs = self._compute_observation(data)
# Get bounce tracking from info
# Get bounce tracking and target heights from info
consecutive_bounces = state.info.get("consecutive_bounces", np.zeros(data.shape[0], dtype=np.int32))
ball_was_upward = state.info.get("ball_was_upward", np.zeros(data.shape[0], dtype=bool))
# Use mean of target_height_range as fallback
default_height = np.mean(self._cfg.target_height_range)
target_heights = state.info.get("target_heights", np.full(data.shape[0], default_height, dtype=np.float32))
# Compute observation with target heights
obs = self._compute_observation(data, target_heights)
# Detect bounces and update consecutive bounce count
current_ball_z = obs[:, 8] # Ball z position
@@ -239,19 +520,42 @@ class BounceBallEnv(NpEnv):
normalized_obs = obs
# Compute reward and termination
reward = self._compute_reward(obs, data, consecutive_bounces)
terminated = self._compute_terminated(obs)
reward, reward_details = self._compute_reward(
obs,
data,
consecutive_bounces,
bounce_detected=bounce_detected,
target_heights=target_heights,
current_actions=state.info.get("current_actions"),
last_actions=state.info.get("last_actions"),
)
terminated = self._compute_terminated(obs, target_heights=target_heights)
state.obs = normalized_obs
state.reward = reward
state.terminated = terminated
# Store reward details for debugging
if self._cfg.store_reward_details:
state.info["Reward"] = reward_details
state.info["target_heights"] = target_heights # Ensure target_heights persists across steps
return state
def reset(self, data: mtx.SceneData) -> tuple:
"""Reset environment to initial state"""
"""Reset environment to initial state with randomized target heights"""
cfg: BounceBallEnvCfg = self._cfg
num_reset = data.shape[0]
# Randomize target heights for the environments being reset
if cfg.randomize_target_height:
min_height, max_height = cfg.target_height_range
new_target_heights = np.random.uniform(min_height, max_height, num_reset).astype(np.float32)
else:
# Use mean of target_height_range when not randomizing
default_height = np.mean(cfg.target_height_range)
new_target_heights = np.full(num_reset, default_height, dtype=np.float32)
# Add noise to initial arm joint positions only (not ball)
arm_noise_pos = np.random.uniform(
-cfg.reset_noise_scale,
@@ -267,42 +571,56 @@ class BounceBallEnv(NpEnv):
# Reset simulation first to get proper DOF structure
data.reset(self._model)
# Get current DOF positions and modify only the arm joints
# Get current DOF positions
current_dof_pos = data.dof_pos
current_dof_vel = data.dof_vel
# === Modify all DOF positions ===
# Set arm joint positions (first 6 DOFs)
current_dof_pos[:, :6] = np.tile(self._init_arm_qpos, (num_reset, 1)) + arm_noise_pos
current_dof_vel[:, :6] = noise_vel[:, :6]
# Set the quaternion part properly (DOFs 9-12 are quaternion w,x,y,z for freejoint)
# The ball has a freejoint which uses quaternion representation
for i in range(num_reset):
# Set quaternion for ball (indices 9-12: w, x, y, z)
current_dof_pos[i, 9:13] = [1.0, 0.0, 0.0, 0.0] # Identity quaternion
data.set_dof_pos(current_dof_pos, self._model)
data.set_dof_vel(current_dof_vel)
# Set ball position in DOF (indices 6-8 for x, y, z positions)
for i in range(num_reset):
ball_noise_pos = np.random.uniform(-0.01, 0.01, 3)
ball_pos = self._ball_init_pos + ball_noise_pos
# Set ball position in DOF coordinates (indices 6-8)
current_dof_pos[i, 6:9] = ball_pos
ball_noise_pos = np.random.uniform(-0.01, 0.01, (num_reset, 3))
ball_pos = self._ball_init_pos + ball_noise_pos
current_dof_pos[:, 6:9] = ball_pos
# Final update to set both ball position and quaternion
# Apply all DOF position changes
data.set_dof_pos(current_dof_pos, self._model)
# Get current DOF velocities
current_dof_vel = data.dof_vel
# === Modify all DOF velocities ===
# Set arm joint velocities (first 6 DOFs)
current_dof_vel[:, :6] = noise_vel[:, :6]
# Apply ball linear velocity in DOF
data.set_dof_vel(current_dof_vel)
# Update target height marker position (thin cylinder disc)
# Marker center aligns with ball top: marker_z = target_height + ball_radius
target_marker_poses = np.tile(self._target_marker_base_pose, (num_reset, 1))
target_marker_poses[:, 2] = new_target_heights + self._ball_radius # Set z position
self._target_marker_body.mocap.set_pose(data, target_marker_poses)
# Update paddle home marker position
paddle_home_marker_poses = np.tile(self._paddle_home_marker_pose, (num_reset, 1))
# Set paddle home marker mocap body pose
self._paddle_home_marker_body.mocap.set_pose(data, paddle_home_marker_poses)
# Initialize info dict with bounce tracking variables
info = {
"consecutive_bounces": np.zeros(num_reset, dtype=np.int32),
"ball_was_upward": np.zeros(num_reset, dtype=bool),
"max_consecutive_bounces": 0,
"target_heights": new_target_heights.copy(), # Return target heights for this reset batch
"current_actions": np.zeros((num_reset, 6), dtype=np.float32),
"last_actions": np.zeros((num_reset, 6), dtype=np.float32),
}
# Compute initial observation
obs = self._compute_observation(data)
# Compute initial observation with target heights
obs = self._compute_observation(data, new_target_heights)
normalized_obs = obs # No normalization for now
return normalized_obs, info

View File

@@ -42,19 +42,110 @@ class BounceBallEnvCfg(EnvCfg):
ball_init_vel: list = None
arm_init_qpos: list = None
# Target positions
target_ball_x: float = 0.58856 # Target x position (m)
target_ball_y: float = 0.0 # Target y position (m)
# Target height for bouncing (configurable parameter)
target_ball_height: float = 0.8 # Default target height in meters
target_height_range: tuple = (0.53, 0.83) # Range for random target height in meters
randomize_target_height: bool = True # Whether to randomize target height on reset
height_tolerance: float = 0.1 # Tolerance for reward calculation
# Paddle behavior constraints
paddle_home_position_z: float = 0.28 # Home position height for paddle (m)
max_paddle_height_deviation: float = 0.1 # Max deviation from home position (m)
encourage_return_home: bool = True # Encourage paddle to return to home position
encourage_impact_velocity: bool = True # Encourage high upward velocity at impact
# Physics constants
gravity: float = 9.81 # Gravity acceleration (m/s^2)
# Reward function parameters
# Position reward
vertical_weight_scale: float = 0.15 # Scale for vertical distance weight
weighted_horizontal_base_scale: float = 0.1 # Base scale for horizontal position
weighted_horizontal_weight_factor: float = 3.0 # Weight factor for vertical proximity
# Out of position penalty
out_of_position_threshold: float = 0.05 # Threshold distance (m)
out_of_position_sharpness: float = 0.03 # Sigmoid sharpness
# Velocity matching
desired_velocity_at_target: float = 0.5 # Target velocity at target height (m/s)
velocity_error_sigma: float = 0.8 # Sigma for velocity error Gaussian
# Height reward
height_error_sigma: float = 0.15 # Sigma for height error Gaussian
height_progress_scale: float = 2.0 # Scale for height progress bonus
height_progress_threshold: float = 0.2 # Minimum height for progress bonus (m)
# Controlled upward velocity
positioning_quality_sigma: float = 0.02 # Sigma for positioning quality
ideal_velocity_min: float = 0.5 # Minimum ideal launch velocity (m/s)
ideal_velocity_max: float = 3.0 # Maximum ideal launch velocity (m/s)
upward_velocity_sigma: float = 0.5 # Sigma for upward velocity quality
upward_mask_scale: float = 0.1 # Scale for upward mask sigmoid
controlled_upward_clip_max: float = 1.5 # Max clip value for controlled upward
# Velocity penalties
excessive_upward_threshold: float = 3.5 # Threshold for excessive upward velocity (m/s)
excessive_upward_sharpness: float = 0.3 # Sigmoid sharpness
downward_velocity_threshold: float = -0.2 # Threshold for downward penalty (m/s)
downward_velocity_scale: float = 0.3 # Scale for downward penalty magnitude
downward_velocity_clip_max: float = 0.5 # Max clip for downward penalty
# Bounce rewards
bounce_positioning_sigma: float = 0.05 # Sigma for bounce positioning quality
bounce_log_scale: float = 0.5 # Scale for logarithmic bounce reward
high_bounce_threshold: float = 2.0 # Threshold for high bounce bonus
high_bounce_sharpness: float = 0.5 # Sigmoid sharpness for high bounce
high_bounce_scale: float = 0.15 # Scale for high bounce bonus
# Paddle alignment
paddle_alignment_sigma: float = 0.03 # Sigma for paddle-ball alignment
vertical_proximity_scale: float = 0.1 # Scale for vertical proximity weight
paddle_alignment_weight_factor: float = 2.0 # Weight factor for vertical proximity
bounce_boost_factor: float = 2.0 # Boost factor when bounce detected
paddle_center_scale: float = 0.3 # Overall scale for paddle center reward
# Home position reward
home_position_sigma: float = 0.05 # Sigma for home position deviation
distance_factor_threshold: float = 0.15 # Threshold for distance factor (m)
distance_factor_sharpness: float = 0.03 # Sigmoid sharpness
distance_factor_scale: float = 0.5 # Scale for distance factor
height_violation_scale: float = 20.0 # Scale for height violation penalty
# Action and velocity penalties
action_penalty_rate: float = 1e-4 # Penalty rate for action changes
joint_vel_penalty_rate: float = 1e-4 # Penalty rate for joint velocities
# Reward weights
weighted_position_weight: float = 2.0
velocity_matching_weight: float = 2.0
height_reward_weight: float = 4.5
height_progress_weight: float = 1.0
controlled_upward_weight: float = 1.5
consecutive_bounces_weight: float = 0.8
high_bounce_weight: float = 0.3
paddle_center_weight: float = 0.6
home_position_weight: float = 1.5
out_of_position_weight: float = 1.0
excessive_upward_weight: float = 1.0
downward_velocity_weight: float = 1.0
height_violation_weight: float = 1.0
# Action scaling parameters
action_scale: list = None
action_bias: list = None
# Debug options
store_reward_details: bool = False # Whether to store detailed reward breakdown in state.info, it's very slow
def __post_init__(self):
if self.ball_init_pos is None:
self.ball_init_pos = [0.58856, 0, 1.27796] # Slightly above paddle (paddle z=0.2803)
self.ball_init_pos = [0.58856, 0, 0.68]
if self.ball_init_vel is None:
self.ball_init_vel = [0.0, 0.0, 0.0]
self.ball_init_vel = [0.0, 0.0, -0.2]
if self.arm_init_qpos is None:
self.arm_init_qpos = [0, 40, 110, 0, -60, 0]

View File

@@ -30,4 +30,4 @@ class CheetahEnvCfg(EnvCfg):
render_spacing: float = 2.0
sim_dt: float = 0.01
ctrl_dt: float = 0.025
run_speed: float = 30.0
run_speed: float = 10.0

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

View File

@@ -0,0 +1,102 @@
# 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__) + "/finger.xml"
spin_model_file = os.path.dirname(__file__) + "/finger_spin.xml"
turn_easy_model_file = os.path.dirname(__file__) + "/finger_turn_easy.xml"
turn_hard_model_file = os.path.dirname(__file__) + "/finger_turn_hard.xml"
@dataclass
class FingerBaseCfg(EnvCfg):
model_file: str = model_file
max_episode_seconds: float = 20.0
sim_dt: float = 0.01
ctrl_dt: float = 0.02
# Task setup
task: str = "spin" # "spin" | "turn"
target_radius: float = 0.07
# Reward thresholds (match dm_control defaults)
spin_velocity_threshold: float = 15.0
# Reward mode
# - "sparse": match dm_control (1 if hinge_velocity <= -threshold else 0)
# - "shaped": dense reward to make training easier
reward_mode: str = "sparse"
shaped_reward_beta: float = 1.0
# Extra shaping for Spin tasks (helps reduce "no contact" failures)
spin_touch_bonus_scale: float = 0.0
spin_touch_bonus_tanh_scale: float = 50.0
spin_approach_reward_scale: float = 0.0
spin_approach_sigma: float = 0.15
# Turn shaping: reward falls linearly to 0 at margin = scale * target_radius
turn_reward_margin_scale: float = 4.0
turn_reward_min_margin: float = 0.0
turn_shaped_reward_beta: float = 1.0
# Turn shaping mode:
# - "linear": clip(1 - max(dist,0)/margin, 0..1)
# - "exp": exp(-max(dist,0)/sigma)
turn_reward_shape: str = "linear"
turn_reward_sigma_scale: float = 1.0
turn_reward_sigma_min: float = 0.05
# Extra shaping for Turn tasks (to reduce jitter and help contact)
turn_touch_bonus_scale: float = 0.05
turn_touch_bonus_tanh_scale: float = 50.0
# Encourage approaching the spinner (helps avoid "no contact" deadlock)
turn_approach_reward_scale: float = 0.3
turn_approach_sigma: float = 0.15
turn_action_l2_penalty_scale: float = 0.002
turn_action_delta_l2_penalty_scale: float = 0.01
# Reset sampling
reset_collision_free_attempts: int = 200
@registry.envcfg("dm-finger-spin")
@dataclass
class FingerSpinCfg(FingerBaseCfg):
model_file: str = spin_model_file
task: str = "spin"
reward_mode: str = "shaped"
spin_approach_reward_scale: float = 0.15
spin_touch_bonus_scale: float = 0.03
@registry.envcfg("dm-finger-turn-easy")
@dataclass
class FingerTurnEasyCfg(FingerBaseCfg):
model_file: str = turn_easy_model_file
task: str = "turn"
target_radius: float = 0.07
reward_mode: str = "shaped"
turn_reward_shape: str = "exp"
@registry.envcfg("dm-finger-turn-hard")
@dataclass
class FingerTurnHardCfg(FingerTurnEasyCfg):
model_file: str = turn_hard_model_file
target_radius: float = 0.03
reward_mode: str = "shaped"
turn_reward_shape: str = "exp"

View File

@@ -0,0 +1,76 @@
<mujoco model="finger">
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<!-- Match cartpole's skybox + floor look -->
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png"/>
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4"/>
</asset>
<option timestep="0.01" cone="elliptic" iterations="200">
<flag gravity="disable"/>
</option>
<default>
<geom solimp="0 0.9 0.01" solref=".02 1"/>
<joint type="hinge" axis="0 -1 0"/>
<motor ctrllimited="true" ctrlrange="-1 1"/>
<default class="finger">
<joint damping="2.5" limited="true"/>
<site type="sphere" size=".03" material="site" group="3"/>
</default>
</default>
<worldbody>
<light name="light" directional="true" diffuse=".6 .6 .6" pos="0 0 2" specular=".3 .3 .3"/>
<geom name="ground" type="plane" pos="0 0 0" size="0 0 0.01" material="motphys-ground"/>
<camera name="cam0" pos="0 -1 .8" xyaxes="1 0 0 0 1 2"/>
<camera name="cam1" pos="0 -1 .4" xyaxes="1 0 0 0 0 1"/>
<body name="proximal" pos="-.2 0 .4" childclass="finger">
<geom name="proximal_decoration" type="cylinder" fromto="0 -.033 0 0 .033 0" size=".034" material="decoration"/>
<joint name="proximal" range="-110 110" ref="-90"/>
<geom name="proximal" type="capsule" material="self" size=".03" fromto="0 0 0 0 0 -.17"/>
<body name="distal" pos="0 0 -.18" childclass="finger">
<joint name="distal" range="-110 110"/>
<geom name="distal" type="capsule" size=".028" material="self" fromto="0 0 0 0 0 -.16" contype="0" conaffinity="0"/>
<geom name="fingertip" type="capsule" size=".03" material="effector" fromto="0 0 -.13 0 0 -.161"/>
<site name="touchtop" pos=".01 0 -.17"/>
<site name="touchbottom" pos="-.01 0 -.17"/>
</body>
</body>
<body name="spinner" pos=".2 0 .4">
<joint name="hinge" frictionloss=".1" damping=".5"/>
<geom name="cap1" type="capsule" size=".04 .09" material="self" pos=".02 0 0"/>
<geom name="cap2" type="capsule" size=".04 .09" material="self" pos="-.02 0 0"/>
<site name="tip" type="sphere" size=".02" pos="0 0 .13" material="target"/>
<geom name="spinner_decoration" type="cylinder" fromto="0 -.045 0 0 .045 0" size=".02" material="decoration"/>
</body>
<site name="target" type="sphere" size=".03" pos="0 0 .4" material="target"/>
</worldbody>
<actuator>
<motor name="proximal" joint="proximal" gear="30"/>
<motor name="distal" joint="distal" gear="15"/>
</actuator>
<!-- All finger observations are functions of sensors. This is useful for finite-differencing. -->
<sensor>
<jointpos name="proximal" joint="proximal"/>
<jointpos name="distal" joint="distal"/>
<jointvel name="proximal_velocity" joint="proximal"/>
<jointvel name="distal_velocity" joint="distal"/>
<jointvel name="hinge_velocity" joint="hinge"/>
<framepos name="tip" objtype="site" objname="tip"/>
<framepos name="target" objtype="site" objname="target"/>
<framepos name="spinner" objtype="xbody" objname="spinner"/>
<touch name="touchtop" site="touchtop"/>
<touch name="touchbottom" site="touchbottom"/>
<framepos name="touchtop_pos" objtype="site" objname="touchtop"/>
<framepos name="touchbottom_pos" objtype="site" objname="touchbottom"/>
</sensor>
</mujoco>

View File

@@ -0,0 +1,431 @@
# 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.finger.cfg import FingerBaseCfg
from motrix_envs.np.env import NpEnv, NpEnvState
def _sanitize_joint_limits(low: np.ndarray, high: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
low = low.copy()
high = high.copy()
low = np.where(np.isfinite(low), low, -np.pi)
high = np.where(np.isfinite(high), high, np.pi)
return low, high
class FingerEnv(NpEnv):
_cfg: FingerBaseCfg
_observation_space: gym.spaces.Box
_action_space: gym.spaces.Box
def __init__(self, cfg: FingerBaseCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
self._cfg = cfg
self._spinner = self._model.get_link("spinner")
self._tip_site = self._model.get_site("tip")
self._target_site = self._model.get_site("target")
self._cap1 = self._model.get_geom("cap1")
self._touchtop_site = self._model.get_site("touchtop")
self._touchbottom_site = self._model.get_site("touchbottom")
self._joint_limit_low, self._joint_limit_high = _sanitize_joint_limits(*self._model.joint_limits)
# Cache joint dof indices
self._prox_qpos_i = self._joint_pos_index("proximal")
self._dist_qpos_i = self._joint_pos_index("distal")
self._hinge_qpos_i = self._joint_pos_index("hinge")
self._hinge_qvel_i = self._joint_vel_index("hinge")
self._prox_qvel_i = self._joint_vel_index("proximal")
self._dist_qvel_i = self._joint_vel_index("distal")
self._target_xyz = np.zeros((num_envs, 3), dtype=np.float32)
self._target_radius = float(cfg.target_radius)
self._spin_vel_threshold = float(cfg.spin_velocity_threshold)
self._init_obs_space()
self._init_action_space()
def _joint_pos_index(self, joint_name: str) -> int:
joint_index = self._model.get_joint_index(joint_name)
return int(self._model.joint_dof_pos_indices[joint_index])
def _joint_vel_index(self, joint_name: str) -> int:
joint_index = self._model.get_joint_index(joint_name)
return int(self._model.joint_dof_vel_indices[joint_index])
def _init_obs_space(self):
raise NotImplementedError
def _init_action_space(self):
low, high = self._model.actuator_ctrl_limits
self._action_space = gym.spaces.Box(low, high, (self._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: np.ndarray, state: NpEnvState) -> NpEnvState:
# Keep track of actions for reward shaping (e.g., smoothness penalties)
if "actions" not in state.info:
state.info["actions"] = np.zeros_like(actions, dtype=np.float32)
if "last_actions" not in state.info:
state.info["last_actions"] = np.zeros_like(actions, dtype=np.float32)
state.info["last_actions"] = state.info["actions"]
state.info["actions"] = actions
state.data.actuator_ctrls = actions
return state
def _touch(self, data: mtx.SceneData) -> np.ndarray:
top = np.asarray(self._model.get_sensor_value("touchtop", data)).reshape(data.shape[0], -1)[:, 0]
bottom = np.asarray(self._model.get_sensor_value("touchbottom", data)).reshape(data.shape[0], -1)[:, 0]
return np.log1p(np.stack([top, bottom], axis=-1))
def _tip_position_xz(self, data: mtx.SceneData) -> np.ndarray:
tip_xyz = self._tip_site.get_position(data)
spinner_xyz = self._spinner.get_position(data)
return (tip_xyz - spinner_xyz)[:, [0, 2]]
def _target_position_xz(self, data: mtx.SceneData) -> np.ndarray:
spinner_xyz = self._spinner.get_position(data)
return (self._target_xyz - spinner_xyz)[:, [0, 2]]
def _dist_to_target(self, data: mtx.SceneData) -> np.ndarray:
# Signed distance to the target surface. Negative means inside.
tip_xyz = self._tip_site.get_position(data)
dist = np.linalg.norm((self._target_xyz - tip_xyz)[:, [0, 2]], axis=-1)
return dist - self._target_radius
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
raise NotImplementedError
def update_state(self, state: NpEnvState) -> NpEnvState:
raise NotImplementedError
def _maybe_init_target_freejoint(self, dof_pos: np.ndarray) -> slice | None:
# Optional freejoint-backed target visualization body (7 qpos: xyz + quat).
try:
self._model.get_geom("target_geom")
target_free_pos = slice(self._model.num_dof_pos - 7, self._model.num_dof_pos)
dof_pos[:, target_free_pos] = np.array([0.0, 0.0, 0.4, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
return target_free_pos
except Exception:
return None
def _reset_collision_free_joint_angles(
self, data: mtx.SceneData, dof_pos: np.ndarray, target_free_pos: slice | None
):
# Randomize joint angles with a collision-free rejection sampler (dm_control-style).
# The MotrixSim joint_limits are per-joint (not per-DOF), so we explicitly fill each DOF.
num = int(data.shape[0])
max_attempts = int(getattr(self._cfg, "reset_collision_free_attempts", 200))
pending = np.ones((num,), dtype=bool)
for _ in range(max_attempts):
if not pending.any():
break
num_pending = int(pending.sum())
for joint_name in ("proximal", "distal"):
j = self._model.get_joint_index(joint_name)
dof_i = self._joint_pos_index(joint_name)
low = float(self._joint_limit_low[j])
high = float(self._joint_limit_high[j])
dof_pos[pending, dof_i] = np.random.uniform(low=low, high=high, size=(num_pending,)).astype(np.float32)
# Sample hinge position (unlimited in model)
dof_pos[pending, self._hinge_qpos_i] = np.random.uniform(
low=-np.pi, high=np.pi, size=(num_pending,)
).astype(np.float32)
if target_free_pos is not None:
dof_pos[:, target_free_pos] = np.array([0.0, 0.0, 0.4, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
data.set_dof_pos(dof_pos, self._model)
data.set_dof_vel(np.zeros((num, self._model.num_dof_vel), dtype=np.float32))
self._model.forward_kinematic(data)
pending = self._model.get_contact_query(data).num_contacts > 0
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
raise NotImplementedError
@registry.env("dm-finger-spin", "np")
class FingerSpinEnv(FingerEnv):
def _init_obs_space(self):
# Match dm_control's observation dict, but flatten into a vector.
# Spin: position(4) + velocity(3) + touch(2) = 9
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32)
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
qpos = data.dof_pos
qvel = data.dof_vel
position = np.concatenate(
[
qpos[:, [self._prox_qpos_i, self._dist_qpos_i]],
self._tip_position_xz(data),
],
axis=-1,
)
velocity = qvel[:, [self._prox_qvel_i, self._dist_qvel_i, self._hinge_qvel_i]]
touch = self._touch(data)
return np.concatenate([position, velocity, touch], axis=-1).astype(np.float32)
def update_state(self, state: NpEnvState) -> NpEnvState:
data = state.data
obs = self._get_obs(data)
terminated = np.isnan(obs).any(axis=-1)
hinge_velocity = data.dof_vel[:, self._hinge_qvel_i]
spin_sparse = (hinge_velocity <= -self._spin_vel_threshold).astype(np.float32)
if self._cfg.reward_mode == "shaped":
# Dense reward to help PPO learn: encourage fast negative hinge velocity.
# Range: [0, 1] roughly, with 1 around reaching the threshold.
spin = np.clip((-hinge_velocity) / self._spin_vel_threshold, 0.0, 1.0).astype(np.float32)
if self._cfg.shaped_reward_beta != 1.0:
spin = np.power(spin, self._cfg.shaped_reward_beta, dtype=np.float32)
else:
spin = spin_sparse
touch_raw = np.zeros((data.shape[0],), dtype=np.float32)
touch_bonus = np.zeros((data.shape[0],), dtype=np.float32)
approach_dist = np.zeros((data.shape[0],), dtype=np.float32)
approach_reward = np.zeros((data.shape[0],), dtype=np.float32)
if self._cfg.reward_mode == "shaped":
if float(getattr(self._cfg, "spin_touch_bonus_scale", 0.0)) > 0.0:
top = np.asarray(self._model.get_sensor_value("touchtop", data)).reshape(data.shape[0], -1)[:, 0]
bottom = np.asarray(self._model.get_sensor_value("touchbottom", data)).reshape(data.shape[0], -1)[:, 0]
touch_raw = (top + bottom).astype(np.float32)
touch_bonus = (
float(self._cfg.spin_touch_bonus_scale)
* np.tanh(touch_raw / float(max(self._cfg.spin_touch_bonus_tanh_scale, 1e-6)))
).astype(np.float32)
if float(getattr(self._cfg, "spin_approach_reward_scale", 0.0)) > 0.0:
spinner_xyz = self._spinner.get_position(data)
top_xyz = self._touchtop_site.get_position(data)
bottom_xyz = self._touchbottom_site.get_position(data)
top_dist = np.linalg.norm((top_xyz - spinner_xyz)[:, [0, 2]], axis=-1)
bottom_dist = np.linalg.norm((bottom_xyz - spinner_xyz)[:, [0, 2]], axis=-1)
approach_dist = np.minimum(top_dist, bottom_dist).astype(np.float32)
sigma = float(max(self._cfg.spin_approach_sigma, 1e-6))
approach_reward = (float(self._cfg.spin_approach_reward_scale) * np.exp(-approach_dist / sigma)).astype(
np.float32
)
spin = np.clip(spin + touch_bonus + approach_reward, 0.0, 1.0).astype(np.float32)
rwd = spin
state.info["Reward"] = {
"hinge_velocity": hinge_velocity.copy(),
"spin": spin.copy(),
"spin_sparse": spin_sparse.copy(),
"touch_raw": touch_raw.copy(),
"touch_bonus": touch_bonus.copy(),
"approach_dist": approach_dist.copy(),
"approach_reward": approach_reward.copy(),
}
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 = int(data.shape[0])
dof_pos = np.zeros((num, self._model.num_dof_pos), dtype=np.float32)
target_free_pos = self._maybe_init_target_freejoint(dof_pos)
self._reset_collision_free_joint_angles(data, dof_pos, target_free_pos)
info: dict = {"Reward": {}}
info["actions"] = np.zeros((num, self._model.num_actuators), dtype=np.float32)
info["last_actions"] = np.zeros((num, self._model.num_actuators), dtype=np.float32)
info["Reward"] = {
"hinge_velocity": np.zeros((num,), dtype=np.float32),
"spin": np.zeros((num,), dtype=np.float32),
"spin_sparse": np.zeros((num,), dtype=np.float32),
"touch_raw": np.zeros((num,), dtype=np.float32),
"touch_bonus": np.zeros((num,), dtype=np.float32),
"approach_dist": np.zeros((num,), dtype=np.float32),
"approach_reward": np.zeros((num,), dtype=np.float32),
}
obs = self._get_obs(data)
return obs, info
@registry.env("dm-finger-turn-easy", "np")
@registry.env("dm-finger-turn-hard", "np")
class FingerTurnEnv(FingerEnv):
def _init_obs_space(self):
# Match dm_control's observation dict, but flatten into a vector.
# Turn: position(4) + velocity(3) + touch(2) + target_position(2) + dist_to_target(1) = 12
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (12,), dtype=np.float32)
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
qpos = data.dof_pos
qvel = data.dof_vel
position = np.concatenate(
[
qpos[:, [self._prox_qpos_i, self._dist_qpos_i]],
self._tip_position_xz(data),
],
axis=-1,
)
velocity = qvel[:, [self._prox_qvel_i, self._dist_qvel_i, self._hinge_qvel_i]]
touch = self._touch(data)
target_position = self._target_position_xz(data)
dist_to_target = self._dist_to_target(data).reshape(data.shape[0], 1)
return np.concatenate([position, velocity, touch, target_position, dist_to_target], axis=-1).astype(np.float32)
def update_state(self, state: NpEnvState) -> NpEnvState:
data = state.data
obs = self._get_obs(data)
terminated = np.isnan(obs).any(axis=-1)
dist_to_target = self._dist_to_target(data)
turn_sparse = (dist_to_target <= 0.0).astype(np.float32)
touch_raw = np.zeros((data.shape[0],), dtype=np.float32)
touch_bonus = np.zeros((data.shape[0],), dtype=np.float32)
approach_dist = np.zeros((data.shape[0],), dtype=np.float32)
approach_reward = np.zeros((data.shape[0],), dtype=np.float32)
action_l2 = np.zeros((data.shape[0],), dtype=np.float32)
action_delta_l2 = np.zeros((data.shape[0],), dtype=np.float32)
if self._cfg.reward_mode == "shaped":
# Encourage approaching the spinner so the agent actually makes contact and can rotate it.
spinner_xyz = self._spinner.get_position(data)
top_xyz = self._touchtop_site.get_position(data)
bottom_xyz = self._touchbottom_site.get_position(data)
top_dist = np.linalg.norm((top_xyz - spinner_xyz)[:, [0, 2]], axis=-1)
bottom_dist = np.linalg.norm((bottom_xyz - spinner_xyz)[:, [0, 2]], axis=-1)
approach_dist = np.minimum(top_dist, bottom_dist).astype(np.float32)
sigma = max(float(self._cfg.turn_approach_sigma), 1e-6)
approach_reward = (self._cfg.turn_approach_reward_scale * np.exp(-approach_dist / sigma)).astype(np.float32)
dist_pos = np.maximum(dist_to_target, 0.0).astype(np.float32)
if getattr(self._cfg, "turn_reward_shape", "linear") == "exp":
sigma = float(
max(self._cfg.turn_reward_sigma_scale * self._target_radius, self._cfg.turn_reward_sigma_min)
)
sigma = max(sigma, 1e-6)
turn = np.exp(-dist_pos / sigma).astype(np.float32)
else:
margin = float(
max(self._cfg.turn_reward_margin_scale * self._target_radius, self._cfg.turn_reward_min_margin)
)
margin = max(margin, 1e-6)
# Dense reward: 1 inside target sphere, decays to 0 at `margin` outside.
turn = np.clip(1.0 - dist_pos / margin, 0.0, 1.0).astype(np.float32)
if self._cfg.turn_shaped_reward_beta != 1.0:
turn = np.power(turn, self._cfg.turn_shaped_reward_beta, dtype=np.float32)
# Encourage making contact (to actually be able to rotate the spinner)
top = np.asarray(self._model.get_sensor_value("touchtop", data)).reshape(data.shape[0], -1)[:, 0]
bottom = np.asarray(self._model.get_sensor_value("touchbottom", data)).reshape(data.shape[0], -1)[:, 0]
touch_raw = (top + bottom).astype(np.float32)
touch_bonus = self._cfg.turn_touch_bonus_scale * np.tanh(touch_raw / self._cfg.turn_touch_bonus_tanh_scale)
# Reduce jitter: penalize large actions and action changes
actions = state.info.get("actions", data.actuator_ctrls).astype(np.float32)
last_actions = state.info.get("last_actions", actions).astype(np.float32)
action_l2 = np.mean(np.square(actions), axis=-1).astype(np.float32)
action_delta_l2 = np.mean(np.square(actions - last_actions), axis=-1).astype(np.float32)
turn = (
turn
+ approach_reward
+ touch_bonus
- self._cfg.turn_action_l2_penalty_scale * action_l2
- self._cfg.turn_action_delta_l2_penalty_scale * action_delta_l2
).astype(np.float32)
turn = np.clip(turn, 0.0, 1.0).astype(np.float32)
else:
turn = turn_sparse
rwd = turn
state.info["Reward"] = {
"dist_to_target": dist_to_target.copy(),
"turn": turn.copy(),
"turn_sparse": turn_sparse.copy(),
"touch_raw": touch_raw.copy(),
"touch_bonus": touch_bonus.copy(),
"approach_dist": approach_dist.copy(),
"approach_reward": approach_reward.copy(),
"action_l2": action_l2.copy(),
"action_delta_l2": action_delta_l2.copy(),
}
state.info["target_info"] = {"positions": self._target_xyz.copy(), "radius": self._target_radius}
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 = int(data.shape[0])
dof_pos = np.zeros((num, self._model.num_dof_pos), dtype=np.float32)
target_free_pos = self._maybe_init_target_freejoint(dof_pos)
self._reset_collision_free_joint_angles(data, dof_pos, target_free_pos)
hinge_xyz = self._spinner.get_position(data)
# Match dm_control: radius = cap1.geom_size.sum() for capsule (radius + half-length).
radius = float(np.sum(self._cap1.size[:2]))
target_angle = np.random.uniform(-np.pi, np.pi, size=(num,))
target_x = hinge_xyz[:, 0] + radius * np.sin(target_angle)
target_z = hinge_xyz[:, 2] + radius * np.cos(target_angle)
self._target_xyz = np.stack([target_x, hinge_xyz[:, 1], target_z], axis=-1).astype(np.float32)
# Best-effort visualization when num_envs == 1 (site position is model-shared).
if self._num_envs == 1:
try:
self._target_site.local_pos = self._target_xyz[0]
self._target_site.size = np.asarray([self._target_radius], dtype=np.float32)
except Exception:
pass
# If we have a freejoint-backed visual target (geom), set its pose in the state.
if target_free_pos is not None:
dof_pos[:, target_free_pos] = np.concatenate(
[
self._target_xyz.astype(np.float32),
np.tile(np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32), (num, 1)),
],
axis=-1,
)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
info: dict = {"Reward": {}}
info["actions"] = np.zeros((num, self._model.num_actuators), dtype=np.float32)
info["last_actions"] = np.zeros((num, self._model.num_actuators), dtype=np.float32)
info["target_info"] = {"positions": self._target_xyz.copy(), "radius": self._target_radius}
info["Reward"] = {
"dist_to_target": np.zeros((num,), dtype=np.float32),
"turn": np.zeros((num,), dtype=np.float32),
"turn_sparse": np.zeros((num,), dtype=np.float32),
}
obs = self._get_obs(data)
return obs, info

View File

@@ -0,0 +1,77 @@
<mujoco model="finger">
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<!-- Match cartpole's skybox + floor look -->
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png"/>
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4"/>
</asset>
<option timestep="0.01" cone="elliptic" iterations="200">
<flag gravity="disable"/>
</option>
<default>
<geom solimp="0 0.9 0.01" solref=".02 1"/>
<joint type="hinge" axis="0 -1 0"/>
<motor ctrllimited="true" ctrlrange="-1 1"/>
<default class="finger">
<joint damping="2.5" limited="true"/>
<site type="sphere" size=".03" material="site" group="3"/>
</default>
</default>
<worldbody>
<light name="light" directional="true" diffuse=".6 .6 .6" pos="0 0 2" specular=".3 .3 .3"/>
<geom name="ground" type="plane" pos="0 0 0" size="0 0 0.01" material="motphys-ground"/>
<camera name="cam0" pos="0 -1 .8" xyaxes="1 0 0 0 1 2"/>
<camera name="cam1" pos="0 -1 .4" xyaxes="1 0 0 0 0 1"/>
<body name="proximal" pos="-.2 0 .4" childclass="finger">
<geom name="proximal_decoration" type="cylinder" fromto="0 -.033 0 0 .033 0" size=".034" material="decoration"/>
<joint name="proximal" range="-110 110" ref="-90"/>
<geom name="proximal" type="capsule" material="self" size=".03" fromto="0 0 0 0 0 -.17"/>
<body name="distal" pos="0 0 -.18" childclass="finger">
<joint name="distal" range="-110 110"/>
<geom name="distal" type="capsule" size=".028" material="self" fromto="0 0 0 0 0 -.16" contype="0" conaffinity="0"/>
<geom name="fingertip" type="capsule" size=".03" material="effector" fromto="0 0 -.13 0 0 -.161"/>
<site name="touchtop" pos=".01 0 -.17"/>
<site name="touchbottom" pos="-.01 0 -.17"/>
</body>
</body>
<body name="spinner" pos=".2 0 .4">
<!-- dm_control's Spin task reduces hinge damping for easier spinning -->
<joint name="hinge" frictionloss=".1" damping=".03"/>
<geom name="cap1" type="capsule" size=".04 .09" material="self" pos=".02 0 0"/>
<geom name="cap2" type="capsule" size=".04 .09" material="self" pos="-.02 0 0"/>
<site name="tip" type="sphere" size=".02" pos="0 0 .13" material="target"/>
<geom name="spinner_decoration" type="cylinder" fromto="0 -.045 0 0 .045 0" size=".02" material="decoration"/>
</body>
<site name="target" type="sphere" size=".03" pos="0 0 .4" material="target"/>
</worldbody>
<actuator>
<motor name="proximal" joint="proximal" gear="30"/>
<motor name="distal" joint="distal" gear="15"/>
</actuator>
<sensor>
<jointpos name="proximal" joint="proximal"/>
<jointpos name="distal" joint="distal"/>
<jointvel name="proximal_velocity" joint="proximal"/>
<jointvel name="distal_velocity" joint="distal"/>
<jointvel name="hinge_velocity" joint="hinge"/>
<framepos name="tip" objtype="site" objname="tip"/>
<framepos name="target" objtype="site" objname="target"/>
<framepos name="spinner" objtype="xbody" objname="spinner"/>
<touch name="touchtop" site="touchtop"/>
<touch name="touchbottom" site="touchbottom"/>
<framepos name="touchtop_pos" objtype="site" objname="touchtop"/>
<framepos name="touchbottom_pos" objtype="site" objname="touchbottom"/>
</sensor>
</mujoco>

View File

@@ -0,0 +1,83 @@
<mujoco model="finger">
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<!-- Match cartpole's skybox + floor look -->
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png"/>
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4"/>
</asset>
<option timestep="0.01" cone="elliptic" iterations="200">
<flag gravity="disable"/>
</option>
<default>
<geom solimp="0 0.9 0.01" solref=".02 1"/>
<joint type="hinge" axis="0 -1 0"/>
<motor ctrllimited="true" ctrlrange="-1 1"/>
<default class="finger">
<joint damping="2.5" limited="true"/>
<site type="sphere" size=".03" material="site" group="3"/>
</default>
</default>
<worldbody>
<light name="light" directional="true" diffuse=".6 .6 .6" pos="0 0 2" specular=".3 .3 .3"/>
<geom name="ground" type="plane" pos="0 0 0" size="0 0 0.01" material="motphys-ground"/>
<camera name="cam0" pos="0 -1 .8" xyaxes="1 0 0 0 1 2"/>
<camera name="cam1" pos="0 -1 .4" xyaxes="1 0 0 0 0 1"/>
<body name="proximal" pos="-.2 0 .4" childclass="finger">
<geom name="proximal_decoration" type="cylinder" fromto="0 -.033 0 0 .033 0" size=".034" material="decoration"/>
<joint name="proximal" range="-110 110" ref="-90"/>
<geom name="proximal" type="capsule" material="self" size=".03" fromto="0 0 0 0 0 -.17"/>
<body name="distal" pos="0 0 -.18" childclass="finger">
<joint name="distal" range="-110 110"/>
<geom name="distal" type="capsule" size=".028" material="self" fromto="0 0 0 0 0 -.16" contype="0" conaffinity="0"/>
<geom name="fingertip" type="capsule" size=".03" material="effector" fromto="0 0 -.13 0 0 -.161"/>
<site name="touchtop" pos=".01 0 -.17"/>
<site name="touchbottom" pos="-.01 0 -.17"/>
</body>
</body>
<body name="spinner" pos=".2 0 .4">
<joint name="hinge" frictionloss=".1" damping=".5"/>
<geom name="cap1" type="capsule" size=".04 .09" material="self" pos=".02 0 0"/>
<geom name="cap2" type="capsule" size=".04 .09" material="self" pos="-.02 0 0"/>
<site name="tip" type="sphere" size=".02" pos="0 0 .13" material="target"/>
<geom name="spinner_decoration" type="cylinder" fromto="0 -.045 0 0 .045 0" size=".02" material="decoration"/>
</body>
<!-- Target as a Site (used for sensors) -->
<site name="target" type="sphere" size=".03" pos="0 0 .4" material="target"/>
<!-- Visible target geometry (MotrixSim renderer may not draw sites) -->
<body name="target_vis" pos="0 0 .4">
<freejoint name="target_free"/>
<geom name="target_geom" type="sphere" size=".07" material="target" contype="0" conaffinity="0"/>
</body>
</worldbody>
<actuator>
<motor name="proximal" joint="proximal" gear="30"/>
<motor name="distal" joint="distal" gear="15"/>
</actuator>
<!-- All finger observations are functions of sensors. This is useful for finite-differencing. -->
<sensor>
<jointpos name="proximal" joint="proximal"/>
<jointpos name="distal" joint="distal"/>
<jointvel name="proximal_velocity" joint="proximal"/>
<jointvel name="distal_velocity" joint="distal"/>
<jointvel name="hinge_velocity" joint="hinge"/>
<framepos name="tip" objtype="site" objname="tip"/>
<framepos name="target" objtype="site" objname="target"/>
<framepos name="spinner" objtype="xbody" objname="spinner"/>
<touch name="touchtop" site="touchtop"/>
<touch name="touchbottom" site="touchbottom"/>
<framepos name="touchtop_pos" objtype="site" objname="touchtop"/>
<framepos name="touchbottom_pos" objtype="site" objname="touchbottom"/>
</sensor>
</mujoco>

View File

@@ -0,0 +1,82 @@
<mujoco model="finger">
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<!-- Match cartpole's skybox + floor look -->
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png"/>
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4"/>
</asset>
<option timestep="0.01" cone="elliptic" iterations="200">
<flag gravity="disable"/>
</option>
<default>
<geom solimp="0 0.9 0.01" solref=".02 1"/>
<joint type="hinge" axis="0 -1 0"/>
<motor ctrllimited="true" ctrlrange="-1 1"/>
<default class="finger">
<joint damping="2.5" limited="true"/>
<site type="sphere" size=".03" material="site" group="3"/>
</default>
</default>
<worldbody>
<light name="light" directional="true" diffuse=".6 .6 .6" pos="0 0 2" specular=".3 .3 .3"/>
<geom name="ground" type="plane" pos="0 0 0" size="0 0 0.01" material="motphys-ground"/>
<camera name="cam0" pos="0 -1 .8" xyaxes="1 0 0 0 1 2"/>
<camera name="cam1" pos="0 -1 .4" xyaxes="1 0 0 0 0 1"/>
<body name="proximal" pos="-.2 0 .4" childclass="finger">
<geom name="proximal_decoration" type="cylinder" fromto="0 -.033 0 0 .033 0" size=".034" material="decoration"/>
<joint name="proximal" range="-110 110" ref="-90"/>
<geom name="proximal" type="capsule" material="self" size=".03" fromto="0 0 0 0 0 -.17"/>
<body name="distal" pos="0 0 -.18" childclass="finger">
<joint name="distal" range="-110 110"/>
<geom name="distal" type="capsule" size=".028" material="self" fromto="0 0 0 0 0 -.16" contype="0" conaffinity="0"/>
<geom name="fingertip" type="capsule" size=".03" material="effector" fromto="0 0 -.13 0 0 -.161"/>
<site name="touchtop" pos=".01 0 -.17"/>
<site name="touchbottom" pos="-.01 0 -.17"/>
</body>
</body>
<body name="spinner" pos=".2 0 .4">
<joint name="hinge" frictionloss=".1" damping=".5"/>
<geom name="cap1" type="capsule" size=".04 .09" material="self" pos=".02 0 0"/>
<geom name="cap2" type="capsule" size=".04 .09" material="self" pos="-.02 0 0"/>
<site name="tip" type="sphere" size=".02" pos="0 0 .13" material="target"/>
<geom name="spinner_decoration" type="cylinder" fromto="0 -.045 0 0 .045 0" size=".02" material="decoration"/>
</body>
<!-- Target as a Site (used for sensors) -->
<site name="target" type="sphere" size=".03" pos="0 0 .4" material="target"/>
<!-- Visible target geometry (MotrixSim renderer may not draw sites) -->
<body name="target_vis" pos="0 0 .4">
<freejoint name="target_free"/>
<geom name="target_geom" type="sphere" size=".03" material="target" contype="0" conaffinity="0"/>
</body>
</worldbody>
<actuator>
<motor name="proximal" joint="proximal" gear="30"/>
<motor name="distal" joint="distal" gear="15"/>
</actuator>
<sensor>
<jointpos name="proximal" joint="proximal"/>
<jointpos name="distal" joint="distal"/>
<jointvel name="proximal_velocity" joint="proximal"/>
<jointvel name="distal_velocity" joint="distal"/>
<jointvel name="hinge_velocity" joint="hinge"/>
<framepos name="tip" objtype="site" objname="tip"/>
<framepos name="target" objtype="site" objname="target"/>
<framepos name="spinner" objtype="xbody" objname="spinner"/>
<touch name="touchtop" site="touchtop"/>
<touch name="touchbottom" site="touchbottom"/>
<framepos name="touchtop_pos" objtype="site" objname="touchtop"/>
<framepos name="touchbottom_pos" objtype="site" objname="touchbottom"/>
</sensor>
</mujoco>

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

View File

@@ -0,0 +1,84 @@
# 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__) + "/humanoid.xml"
@dataclass
class InitStateConfig:
reset_height_factor: float = 0.95
reset_qvel_range: float = 0.01
reset_actuator_range: float = 0.02
hip_yaw_range: tuple[float, float] = (-15.0, 15.0)
hip_roll_range: tuple[float, float] = (-12.0, 12.0)
hip_pitch_range: tuple[float, float] = (-12.0, 12.0)
symmetric_leg_pairs: list[tuple[int, int, tuple[float, float]]] = field(
default_factory=lambda: [
(10, 16, (-18.0, 2.0)),
(11, 17, (-25.0, 20.0)),
(12, 18, (-70.0, 5.0)),
(13, 19, (-45.0, -25.0)),
(14, 20, (-40.0, 0.0)),
(15, 21, (-25.0, 25.0)),
]
)
symmetric_arm_pairs: list[tuple[int, int]] = field(
default_factory=lambda: [
(22, 25),
(23, 26),
(24, 27),
]
)
arm_margin_factor: float = 0.1
@dataclass
class TerminationConfig:
head_height_factor: float = 0.5
torso_upright_threshold: float = 0.2
extreme_vel_threshold: float = 200.0
@registry.envcfg("dm-humanoid-walk")
@dataclass
class HumanoidWalkCfg(EnvCfg):
model_file: str = model_file
max_episode_seconds: float = 25.0
sim_dt: float = 0.01
ctrl_dt: float = 0.01
move_speed: float = 1.0
stand_height: float = 1.4
init_state: InitStateConfig = field(default_factory=InitStateConfig)
termination_config: TerminationConfig = field(default_factory=TerminationConfig)
@registry.envcfg("dm-humanoid-stand")
@dataclass
class HumanoidStandCfg(HumanoidWalkCfg):
move_speed: float = 0.0
@registry.envcfg("dm-humanoid-run")
@dataclass
class HumanoidRunCfg(HumanoidWalkCfg):
move_speed: float = 10.0

View File

@@ -0,0 +1,207 @@
<mujoco model="humanoid">
<include file="../../common/skybox.xml"/>
<include file="../../common/visual.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<texture name="motphys-ground" type="2d" file="../../common/motphys-ground.png" />
<material name="motphys-ground" texture="motphys-ground" texuniform="true"
texrepeat="0.4 0.4" />
</asset>
<statistic extent="2" center="0 0 1"/>
<option timestep=".005"/>
<default>
<motor ctrlrange="-1 1" ctrllimited="true"/>
<default class="body">
<geom type="capsule" condim="1" friction=".7" solimp=".9 .99 .003" solref=".015 1" material="self"/>
<joint type="hinge" damping=".2" stiffness="1" armature=".01" limited="true" solimplimit="0 .99 .01"/>
<default class="big_joint">
<joint damping="5" stiffness="10"/>
<default class="big_stiff_joint">
<joint stiffness="20"/>
</default>
</default>
<site size=".04" group="3"/>
<default class="force-torque">
<site type="box" size=".01 .01 .02" rgba="1 0 0 1" />
</default>
<default class="touch">
<site type="capsule" rgba="0 0 1 .3"/>
</default>
</default>
</default>
<worldbody>
<geom name="floor" size="0 0 0.01" type="plane" material="motphys-ground" contype="1"
conaffinity="1" priority="1" friction="0.6" condim="3" />
<body name="torso" pos="0 0 1.5" childclass="body">
<light name="sun" pos="0 0 8" dir="0 0 -1" directional="true" ambient="0.4 0.4 0.4" diffuse="0.8 0.8 0.8" specular="0.1 0.1 0.1"/>
<camera name="back" pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="trackcom"/>
<camera name="side" pos="0 -3 1" xyaxes="1 0 0 0 1 2" mode="trackcom"/>
<freejoint name="root"/>
<site name="root" class="force-torque"/>
<geom name="torso" fromto="0 -.07 0 0 .07 0" size=".07"/>
<geom name="upper_waist" fromto="-.01 -.06 -.12 -.01 .06 -.12" size=".06"/>
<site name="torso" class="touch" type="box" pos="0 0 -.05" size=".075 .14 .13"/>
<body name="head" pos="0 0 .19">
<geom name="head" type="sphere" size=".09"/>
<site name="head" class="touch" type="sphere" size=".091"/>
<camera name="egocentric" pos=".09 0 0" xyaxes="0 -1 0 .1 0 1" fovy="80"/>
</body>
<body name="lower_waist" pos="-.01 0 -.260" quat="1.000 0 -.002 0">
<geom name="lower_waist" fromto="0 -.06 0 0 .06 0" size=".06"/>
<site name="lower_waist" class="touch" size=".061 .06" zaxis="0 1 0"/>
<joint name="abdomen_z" pos="0 0 .065" axis="0 0 1" range="-45 45" class="big_stiff_joint"/>
<joint name="abdomen_y" pos="0 0 .065" axis="0 1 0" range="-75 30" class="big_joint"/>
<body name="pelvis" pos="0 0 -.165" quat="1.000 0 -.002 0">
<joint name="abdomen_x" pos="0 0 .1" axis="1 0 0" range="-35 35" class="big_joint"/>
<geom name="butt" fromto="-.02 -.07 0 -.02 .07 0" size=".09"/>
<site name="butt" class="touch" size=".091 .07" pos="-.02 0 0" zaxis="0 1 0"/>
<body name="right_thigh" pos="0 -.1 -.04">
<site name="right_hip" class="force-torque"/>
<joint name="right_hip_x" axis="1 0 0" range="-25 5" class="big_joint"/>
<joint name="right_hip_z" axis="0 0 1" range="-60 35" class="big_joint"/>
<joint name="right_hip_y" axis="0 1 0" range="-120 20" class="big_stiff_joint"/>
<geom name="right_thigh" fromto="0 0 0 0 .01 -.34" size=".06"/>
<site name="right_thigh" class="touch" pos="0 .005 -.17" size=".061 .17" zaxis="0 -1 34"/>
<body name="right_shin" pos="0 .01 -.403">
<site name="right_knee" class="force-torque" pos="0 0 .02"/>
<joint name="right_knee" pos="0 0 .02" axis="0 -1 0" range="-160 2"/>
<geom name="right_shin" fromto="0 0 0 0 0 -.3" size=".049"/>
<site name="right_shin" class="touch" pos="0 0 -.15" size=".05 .15"/>
<body name="right_foot" pos="0 0 -.39">
<site name="right_ankle" class="force-torque"/>
<joint name="right_ankle_y" pos="0 0 .08" axis="0 1 0" range="-50 50" stiffness="6"/>
<joint name="right_ankle_x" pos="0 0 .04" axis="1 0 .5" range="-50 50" stiffness="3"/>
<geom name="right_right_foot" fromto="-.07 -.02 0 .14 -.04 0" size=".027"/>
<geom name="left_right_foot" fromto="-.07 0 0 .14 .02 0" size=".027"/>
<site name="right_right_foot" class="touch" pos=".035 -.03 0" size=".03 .11" zaxis="21 -2 0"/>
<site name="left_right_foot" class="touch" pos=".035 .01 0" size=".03 .11" zaxis="21 2 0"/>
</body>
</body>
</body>
<body name="left_thigh" pos="0 .1 -.04">
<site name="left_hip" class="force-torque"/>
<joint name="left_hip_x" axis="-1 0 0" range="-25 5" class="big_joint"/>
<joint name="left_hip_z" axis="0 0 -1" range="-60 35" class="big_joint"/>
<joint name="left_hip_y" axis="0 1 0" range="-120 20" class="big_stiff_joint"/>
<geom name="left_thigh" fromto="0 0 0 0 -.01 -.34" size=".06"/>
<site name="left_thigh" class="touch" pos="0 -.005 -.17" size=".061 .17" zaxis="0 1 34"/>
<body name="left_shin" pos="0 -.01 -.403">
<site name="left_knee" class="force-torque" pos="0 0 .02"/>
<joint name="left_knee" pos="0 0 .02" axis="0 -1 0" range="-160 2"/>
<geom name="left_shin" fromto="0 0 0 0 0 -.3" size=".049"/>
<site name="left_shin" class="touch" pos="0 0 -.15" size=".05 .15"/>
<body name="left_foot" pos="0 0 -.39">
<site name="left_ankle" class="force-torque"/>
<joint name="left_ankle_y" pos="0 0 .08" axis="0 1 0" range="-50 50" stiffness="6"/>
<joint name="left_ankle_x" pos="0 0 .04" axis="1 0 .5" range="-50 50" stiffness="3"/>
<geom name="left_left_foot" fromto="-.07 .02 0 .14 .04 0" size=".027"/>
<geom name="right_left_foot" fromto="-.07 0 0 .14 -.02 0" size=".027"/>
<site name="right_left_foot" class="touch" pos=".035 -.01 0" size=".03 .11" zaxis="21 -2 0"/>
<site name="left_left_foot" class="touch" pos=".035 .03 0" size=".03 .11" zaxis="21 2 0"/>
</body>
</body>
</body>
</body>
</body>
<body name="right_upper_arm" pos="0 -.17 .06">
<joint name="right_shoulder1" axis="2 1 1" range="-85 60"/>
<joint name="right_shoulder2" axis="0 -1 1" range="-85 60"/>
<geom name="right_upper_arm" fromto="0 0 0 .16 -.16 -.16" size=".04 .16"/>
<site name="right_upper_arm" class="touch" pos=".08 -.08 -.08" size=".041 .14" zaxis="1 -1 -1"/>
<body name="right_lower_arm" pos=".18 -.18 -.18">
<joint name="right_elbow" axis="0 -1 1" range="-90 50" stiffness="0"/>
<geom name="right_lower_arm" fromto=".01 .01 .01 .17 .17 .17" size=".031"/>
<site name="right_lower_arm" class="touch" pos=".09 .09 .09" size=".032 .14" zaxis="1 1 1"/>
<body name="right_hand" pos=".18 .18 .18">
<geom name="right_hand" type="sphere" size=".04"/>
<site name="right_hand" class="touch" type="sphere" size=".041"/>
</body>
</body>
</body>
<body name="left_upper_arm" pos="0 .17 .06">
<joint name="left_shoulder1" axis="2 -1 1" range="-60 85"/>
<joint name="left_shoulder2" axis="0 1 1" range="-60 85"/>
<geom name="left_upper_arm" fromto="0 0 0 .16 .16 -.16" size=".04 .16"/>
<site name="left_upper_arm" class="touch" pos=".08 .08 -.08" size=".041 .14" zaxis="1 1 -1"/>
<body name="left_lower_arm" pos=".18 .18 -.18">
<joint name="left_elbow" axis="0 -1 -1" range="-90 50" stiffness="0"/>
<geom name="left_lower_arm" fromto=".01 -.01 .01 .17 -.17 .17" size=".031"/>
<site name="left_lower_arm" class="touch" pos=".09 -.09 .09" size=".032 .14" zaxis="1 -1 1"/>
<body name="left_hand" pos=".18 -.18 .18">
<geom name="left_hand" type="sphere" size=".04"/>
<site name="left_hand" class="touch" type="sphere" size=".041"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor name="abdomen_y" gear="40" joint="abdomen_y"/>
<motor name="abdomen_z" gear="40" joint="abdomen_z"/>
<motor name="abdomen_x" gear="40" joint="abdomen_x"/>
<motor name="right_hip_x" gear="40" joint="right_hip_x"/>
<motor name="right_hip_z" gear="40" joint="right_hip_z"/>
<motor name="right_hip_y" gear="120" joint="right_hip_y"/>
<motor name="right_knee" gear="80" joint="right_knee"/>
<motor name="right_ankle_x" gear="20" joint="right_ankle_x"/>
<motor name="right_ankle_y" gear="20" joint="right_ankle_y"/>
<motor name="left_hip_x" gear="40" joint="left_hip_x"/>
<motor name="left_hip_z" gear="40" joint="left_hip_z"/>
<motor name="left_hip_y" gear="120" joint="left_hip_y"/>
<motor name="left_knee" gear="80" joint="left_knee"/>
<motor name="left_ankle_x" gear="20" joint="left_ankle_x"/>
<motor name="left_ankle_y" gear="20" joint="left_ankle_y"/>
<motor name="right_shoulder1" gear="20" joint="right_shoulder1"/>
<motor name="right_shoulder2" gear="20" joint="right_shoulder2"/>
<motor name="right_elbow" gear="40" joint="right_elbow"/>
<motor name="left_shoulder1" gear="20" joint="left_shoulder1"/>
<motor name="left_shoulder2" gear="20" joint="left_shoulder2"/>
<motor name="left_elbow" gear="40" joint="left_elbow"/>
</actuator>
<sensor>
<subtreelinvel name="torso_subtreelinvel" body="torso"/>
<accelerometer name="torso_accel" site="root"/>
<velocimeter name="torso_vel" site="root"/>
<gyro name="torso_gyro" site="root"/>
<force name="left_ankle_force" site="left_ankle"/>
<force name="right_ankle_force" site="right_ankle"/>
<force name="left_knee_force" site="left_knee"/>
<force name="right_knee_force" site="right_knee"/>
<force name="left_hip_force" site="left_hip"/>
<force name="right_hip_force" site="right_hip"/>
<torque name="left_ankle_torque" site="left_ankle"/>
<torque name="right_ankle_torque" site="right_ankle"/>
<torque name="left_knee_torque" site="left_knee"/>
<torque name="right_knee_torque" site="right_knee"/>
<torque name="left_hip_torque" site="left_hip"/>
<torque name="right_hip_torque" site="right_hip"/>
<touch name="torso_touch" site="torso"/>
<touch name="head_touch" site="head"/>
<touch name="lower_waist_touch" site="lower_waist"/>
<touch name="butt_touch" site="butt"/>
<touch name="right_thigh_touch" site="right_thigh"/>
<touch name="right_shin_touch" site="right_shin"/>
<touch name="right_right_foot_touch" site="right_right_foot"/>
<touch name="left_right_foot_touch" site="left_right_foot"/>
<touch name="left_thigh_touch" site="left_thigh"/>
<touch name="left_shin_touch" site="left_shin"/>
<touch name="right_left_foot_touch" site="right_left_foot"/>
<touch name="left_left_foot_touch" site="left_left_foot"/>
<touch name="right_upper_arm_touch" site="right_upper_arm"/>
<touch name="right_lower_arm_touch" site="right_lower_arm"/>
<touch name="right_hand_touch" site="right_hand"/>
<touch name="left_upper_arm_touch" site="left_upper_arm"/>
<touch name="left_lower_arm_touch" site="left_lower_arm"/>
<touch name="left_hand_touch" site="left_hand"/>
</sensor>
</mujoco>

View File

@@ -0,0 +1,550 @@
# 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.humanoid.cfg import HumanoidWalkCfg
from motrix_envs.np import reward
from motrix_envs.np.env import NpEnv, NpEnvState
@registry.env("dm-humanoid-stand", "np")
@registry.env("dm-humanoid-walk", "np")
@registry.env("dm-humanoid-run", "np")
class Humanoid3DEnv(NpEnv):
_observation_space: gym.spaces.Box
_action_space: gym.spaces.Box
def __init__(self, cfg: HumanoidWalkCfg, num_envs=1):
super().__init__(cfg, num_envs)
self._init_obs_space()
self._init_action_space()
self._torso = self._model.get_link("torso")
self._head = self._model.get_link("head")
self._pelvis = self._model.get_link("pelvis")
self._left_hand = self._model.get_link("left_hand")
self._right_hand = self._model.get_link("right_hand")
self._left_foot = self._model.get_link("left_foot")
self._right_foot = self._model.get_link("right_foot")
self._move_speed = float(cfg.move_speed)
self._stand_height = float(cfg.stand_height)
self._target_direction = np.array([1.0, 0.0, 0.0], dtype=np.float32)
self._target_direction_xy = self._target_direction[:2].copy()
self._qpos_low, self._qpos_high = self._build_qpos_limits(self._model)
self._cache_derived_constants(cfg)
self._init_joint_randomization_config(cfg)
def _build_qpos_limits(self, model) -> tuple[np.ndarray, np.ndarray]:
num_dof_pos = int(model.num_dof_pos)
jl = np.asarray(model.joint_limits, dtype=np.float32)
if jl.ndim != 2 or jl.shape[0] != 2:
low = np.full((num_dof_pos,), -np.inf, dtype=np.float32)
high = np.full((num_dof_pos,), np.inf, dtype=np.float32)
return low, high
k = int(jl.shape[1])
low = np.full((num_dof_pos,), -np.inf, dtype=np.float32)
high = np.full((num_dof_pos,), np.inf, dtype=np.float32)
m = min(k, num_dof_pos)
low[:m] = jl[0, :m]
high[:m] = jl[1, :m]
return low, high
def _cache_derived_constants(self, cfg: HumanoidWalkCfg) -> None:
t_cfg = cfg.termination_config
self._head_height_min = self._stand_height * 0.95
self._pelvis_height_min = 0.6 * self._stand_height
self._pelvis_height_margin = 0.6 * self._stand_height
self._term_head_height_min = float(t_cfg.head_height_factor) * self._stand_height
self._term_torso_upright_threshold = float(t_cfg.torso_upright_threshold)
self._term_extreme_vel_threshold = float(t_cfg.extreme_vel_threshold)
def _init_obs_space(self):
model = self._model
num_joint_angles = model.num_dof_pos - 7
num_head_height = 1
num_extremities = 12
num_torso_vertical = 3
num_com_vel = 3
num_qvel = model.num_dof_vel
num_target_local = 3
num_obs = (
num_joint_angles
+ num_head_height
+ num_extremities
+ num_torso_vertical
+ num_com_vel
+ num_qvel
+ num_target_local
)
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(
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: np.ndarray, state: NpEnvState) -> NpEnvState:
state.data.actuator_ctrls = actions
return state
def update_state(self, state: NpEnvState) -> NpEnvState:
state = self.update_observation(state)
state = self.update_terminated(state)
state = self.update_reward(state)
return state
def update_observation(self, state: NpEnvState) -> NpEnvState:
data = state.data
obs = self._get_obs(data)
return state.replace(obs=obs)
def update_terminated(self, state: NpEnvState) -> NpEnvState:
data = state.data
head_height = self._get_head_height(data)
torso_upright = self._get_torso_upright(data)
terminated = self._compute_terminated(data, head_height, torso_upright)
return state.replace(
terminated=terminated,
)
def update_reward(self, state: NpEnvState) -> NpEnvState:
data = state.data
terminated = state.terminated
head_height = self._get_head_height(data)
pelvis_height = self._get_pelvis_height(data)
torso_upright = self._get_torso_upright(data)
rwd, reward_components = self._compute_reward(data, head_height, torso_upright, pelvis_height)
rwd, reward_components = self._apply_termination_mask(terminated, rwd, reward_components)
state.info["Reward"] = reward_components
return state.replace(reward=rwd)
def _apply_termination_mask(
self,
terminated: np.ndarray,
rwd: np.ndarray,
reward_components: dict,
) -> tuple[np.ndarray, dict]:
rwd = np.where(terminated, 0.0, rwd).astype(np.float32)
for k, v in reward_components.items():
reward_components[k] = np.where(terminated, 0.0, v).astype(np.float32)
return rwd, reward_components
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
data.reset(self._model)
self._randomize_joints_inplace(data)
obs = self._get_obs(data)
return obs, {}
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
joint_angles = np.asarray(data.dof_pos[:, 7:], dtype=np.float32)
head_height = self._get_head_height(data).astype(np.float32)[:, None]
extremities = self._get_extremities(data).astype(np.float32)
torso_rot = self._torso.get_rotation_mat(data)
torso_vertical = np.asarray(torso_rot[:, 2, :], dtype=np.float32)
com_vel = np.asarray(self._model.get_sensor_value("torso_subtreelinvel", data), dtype=np.float32)
qvel = np.asarray(data.dof_vel, dtype=np.float32)
target_direction_local = self._get_target_direction_local(data).astype(np.float32)
obs = np.concatenate(
[joint_angles, head_height, extremities, torso_vertical, com_vel, qvel, target_direction_local], axis=-1
)
return obs
def _get_head_height(self, data: mtx.SceneData) -> np.ndarray:
return np.asarray(self._head.get_position(data)[:, 2], dtype=np.float32)
def _get_pelvis_height(self, data: mtx.SceneData) -> np.ndarray:
return np.asarray(self._pelvis.get_position(data)[:, 2], dtype=np.float32)
def _get_torso_upright(self, data: mtx.SceneData) -> np.ndarray:
torso_rot = self._torso.get_rotation_mat(data)
return np.asarray(torso_rot[:, 2, 2], dtype=np.float32)
def _get_extremities(self, data: mtx.SceneData) -> np.ndarray:
torso_rot = self._torso.get_rotation_mat(data)
torso_pos = self._torso.get_position(data)
parts = [
self._left_hand.get_position(data),
self._left_foot.get_position(data),
self._right_hand.get_position(data),
self._right_foot.get_position(data),
]
out = []
torso_rot_f32 = np.asarray(torso_rot, dtype=np.float32)
torso_pos_f32 = np.asarray(torso_pos, dtype=np.float32)
for p in parts:
torso_to_limb = np.asarray(p, dtype=np.float32) - torso_pos_f32
v_body = np.einsum("ni,nij->nj", torso_to_limb, torso_rot_f32)
out.append(v_body)
return np.concatenate(out, axis=-1)
def _get_target_direction_local(self, data: mtx.SceneData) -> np.ndarray:
n = int(data.shape[0])
torso_rot = self._torso.get_rotation_mat(data)
torso_rot_f32 = np.asarray(torso_rot, dtype=np.float32)
target_world = np.ones((n, 3), dtype=np.float32) * self._target_direction[None, :]
target_local = np.einsum("ni,nij->nj", target_world, torso_rot_f32)
return target_local
def _compute_reward(
self,
data: mtx.SceneData,
head_height: np.ndarray,
torso_upright: np.ndarray,
pelvis_height: np.ndarray,
) -> tuple[np.ndarray, dict]:
posture_reward = self._compute_posture_reward(head_height, torso_upright, pelvis_height)
speed_reward, energy_reward = self._compute_speed_and_energy_reward(data)
gait_reward = self._compute_gait_reward(data)
rwd = (posture_reward * speed_reward * energy_reward * gait_reward).astype(np.float32)
comps = {
"energy": energy_reward.astype(np.float32),
"speed": speed_reward.astype(np.float32),
"posture": posture_reward.astype(np.float32),
"gait": gait_reward.astype(np.float32),
}
return rwd, comps
def _compute_posture_reward(
self,
head_height: np.ndarray,
torso_upright: np.ndarray,
pelvis_height: np.ndarray,
) -> np.ndarray:
stand_reward = (
reward.tolerance(
head_height,
bounds=(self._head_height_min, float("inf")),
margin=0.5,
)
.astype(np.float32)
.flatten()
)
upright_reward = (
reward.tolerance(
torso_upright,
bounds=(0.9, float("inf")),
sigmoid="linear",
margin=0.9,
)
.astype(np.float32)
.flatten()
)
pelvis_height_reward = (
reward.tolerance(
pelvis_height,
bounds=(self._pelvis_height_min, float("inf")),
sigmoid="linear",
margin=self._pelvis_height_margin,
)
.astype(np.float32)
.flatten()
)
return (stand_reward * upright_reward * pelvis_height_reward).astype(np.float32)
def _compute_speed_and_energy_reward(
self,
data: mtx.SceneData,
) -> tuple[np.ndarray, np.ndarray]:
target_dir_xy = self._target_direction_xy
ctrls = np.asarray(data.actuator_ctrls, dtype=np.float32)
com_vel = np.asarray(self._model.get_sensor_value("torso_subtreelinvel", data), dtype=np.float32)
if self._move_speed <= 0.0:
energy_reward = np.exp(-1.0 * np.mean(np.square(ctrls), axis=-1)).astype(np.float32)
actual_speed = np.linalg.norm(com_vel[:, :2], axis=-1).astype(np.float32)
speed_reward = (
reward.tolerance(
actual_speed,
bounds=(self._move_speed, self._move_speed),
margin=1.0,
value_at_margin=0.01,
)
.astype(np.float32)
.flatten()
)
elif self._move_speed <= 3.0:
energy_reward = np.exp(-0.5 * np.mean(np.square(ctrls), axis=-1)).astype(np.float32)
actual_speed = np.sum(com_vel[:, :2] * target_dir_xy, axis=-1).astype(np.float32)
speed_reward = (
reward.tolerance(
actual_speed,
bounds=(self._move_speed, self._move_speed),
margin=self._move_speed,
value_at_margin=0.0,
sigmoid="linear",
)
.astype(np.float32)
.flatten()
)
else:
energy_reward = np.exp(-0.3 * np.mean(np.square(ctrls), axis=-1)).astype(np.float32)
actual_speed = np.sum(com_vel[:, :2] * target_dir_xy, axis=-1).astype(np.float32)
speed_reward = (
reward.tolerance(
actual_speed,
bounds=(self._move_speed, float("inf")),
margin=self._move_speed,
value_at_margin=0.0,
sigmoid="linear",
)
.astype(np.float32)
.flatten()
)
return speed_reward, energy_reward
def _compute_heading_reward(
self,
forward_vec: np.ndarray,
target_dir: np.ndarray,
bounds,
margin,
) -> np.ndarray:
dot = np.sum(forward_vec * target_dir, axis=-1)
return (
reward.tolerance(
dot,
bounds=bounds,
margin=margin,
value_at_margin=0.0,
sigmoid="linear",
)
.astype(np.float32)
.flatten()
)
def _compute_gait_reward(self, data: mtx.SceneData) -> np.ndarray:
target_dir = self._target_direction
torso_rot = self._torso.get_rotation_mat(data)
head_rot = self._head.get_rotation_mat(data)
pelvis_rot = self._pelvis.get_rotation_mat(data)
torso_forward = np.asarray(torso_rot[:, 0, 0:3], dtype=np.float32)
torso_heading_reward = self._compute_heading_reward(torso_forward, target_dir, bounds=(0.9, 1.0), margin=0.3)
head_forward = np.asarray(head_rot[:, 0, 0:3], dtype=np.float32)
head_heading_reward = self._compute_heading_reward(head_forward, target_dir, bounds=(0.9, 1.0), margin=0.3)
pelvis_forward = np.asarray(pelvis_rot[:, 0, 0:3], dtype=np.float32)
pelvis_yaw_reward = self._compute_heading_reward(pelvis_forward, target_dir, bounds=(0.9, 1.0), margin=0.3)
pelvis_up = np.asarray(pelvis_rot[:, 2, 2], dtype=np.float32)
pelvis_level_reward = (
reward.tolerance(
pelvis_up,
bounds=(0.9, 1.0),
margin=0.3,
sigmoid="linear",
value_at_margin=0.0,
)
.astype(np.float32)
.flatten()
)
left_foot_pos = self._left_foot.get_position(data)
right_foot_pos = self._right_foot.get_position(data)
max_foot_h = np.maximum(
np.asarray(left_foot_pos[:, 2], dtype=np.float32),
np.asarray(right_foot_pos[:, 2], dtype=np.float32),
)
feet_height_reward = (
reward.tolerance(
max_foot_h,
bounds=(0.0, 0.3),
margin=0.5,
sigmoid="quadratic",
value_at_margin=0.0,
)
.astype(np.float32)
.flatten()
)
return (
torso_heading_reward * head_heading_reward * pelvis_yaw_reward * pelvis_level_reward * feet_height_reward
).astype(np.float32)
def _compute_terminated(
self,
data: mtx.SceneData,
head_height: np.ndarray,
torso_upright: np.ndarray,
) -> np.ndarray:
qpos = np.asarray(data.dof_pos, dtype=np.float32)
qvel = np.asarray(data.dof_vel, dtype=np.float32)
bad = ~np.isfinite(qpos).all(axis=-1) | ~np.isfinite(qvel).all(axis=-1)
too_low = head_height < self._term_head_height_min
too_tilted = torso_upright < self._term_torso_upright_threshold
extreme_vel = np.abs(qvel).max(axis=-1) > self._term_extreme_vel_threshold
return (bad | too_low | too_tilted | extreme_vel).astype(bool)
def _init_joint_randomization_config(self, cfg: HumanoidWalkCfg) -> None:
init_cfg = cfg.init_state
self._reset_height = self._stand_height * init_cfg.reset_height_factor
self._reset_qvel_range = init_cfg.reset_qvel_range
self._reset_actuator_range = init_cfg.reset_actuator_range
self._hip_yaw_range = tuple(np.deg2rad(x) for x in init_cfg.hip_yaw_range)
self._hip_roll_range = tuple(np.deg2rad(x) for x in init_cfg.hip_roll_range)
self._hip_pitch_range = tuple(np.deg2rad(x) for x in init_cfg.hip_pitch_range)
self._symmetric_leg_pairs_rad = [
(left_idx, right_idx, tuple(np.deg2rad(x) for x in deg_range))
for left_idx, right_idx, deg_range in init_cfg.symmetric_leg_pairs
]
self._symmetric_arm_pairs = init_cfg.symmetric_arm_pairs
self._arm_margin_factor = init_cfg.arm_margin_factor
self._symmetric_arm_used_indices = set()
for left_idx, right_idx in self._symmetric_arm_pairs:
self._symmetric_arm_used_indices.add(left_idx)
self._symmetric_arm_used_indices.add(right_idx)
def _randomize_joints_inplace(self, data: mtx.SceneData) -> None:
# qpos layout (humanoid.xml): 0-6 free (x,y,z,qw,qx,qy,qz), 7=abdomen_z, 8=abdomen_y, 9=abdomen_x,
# 10-15 right leg (hip_x,z,y, knee, ankle_y,x), 16-21 left leg, 22-24 right arm, 25-27 left arm (num_dof_pos=28)
model = self._model
n = int(data.shape[0])
num_dof_pos = int(model.num_dof_pos)
num_dof_vel = int(model.num_dof_vel)
num_actuators = int(model.num_actuators)
low, high = self._qpos_low, self._qpos_high
qpos = np.zeros((n, num_dof_pos), dtype=np.float32)
qpos[:, 2] = self._reset_height
qpos[:, 3] = 1.0
# qpos 7=abdomen_z (yaw), 8=abdomen_y (pitch), 9=abdomen_x (roll) per humanoid.xml
qpos[:, 7] = np.random.uniform(self._hip_yaw_range[0], self._hip_yaw_range[1], size=(n,))
qpos[:, 8] = np.random.uniform(self._hip_pitch_range[0], self._hip_pitch_range[1], size=(n,))
qpos[:, 9] = np.random.uniform(self._hip_roll_range[0], self._hip_roll_range[1], size=(n,))
self._randomize_symmetric_legs(qpos, n, num_dof_pos, low, high)
self._randomize_symmetric_arms(qpos, n, num_dof_pos, low, high)
self._randomize_remaining_joints(qpos, n, num_dof_pos, low, high)
qvel = np.random.uniform(-self._reset_qvel_range, self._reset_qvel_range, size=(n, num_dof_vel)).astype(
np.float32
)
actuator_ctrls = np.random.uniform(
-self._reset_actuator_range, self._reset_actuator_range, size=(n, num_actuators)
).astype(np.float32)
qpos_set = qpos.copy()
qpos_set[:, 3:7] = np.concatenate([qpos[:, 4:7], qpos[:, 3:4]], axis=1)
data.set_dof_pos(qpos_set, self._model)
data.set_dof_vel(qvel)
data.actuator_ctrls[:] = actuator_ctrls
self._model.forward_kinematic(data)
def _randomize_symmetric_legs(
self, qpos: np.ndarray, n: int, num_dof_pos: int, low: np.ndarray, high: np.ndarray
) -> None:
for left_idx, right_idx, (min_rad, max_rad) in self._symmetric_leg_pairs_rad:
if left_idx < num_dof_pos:
qpos[:, left_idx] = np.random.uniform(
np.clip(min_rad, low[left_idx], high[left_idx]),
np.clip(max_rad, low[left_idx], high[left_idx]),
size=(n,),
)
if right_idx < num_dof_pos:
right_min_rad = -max_rad
right_max_rad = -min_rad
qpos[:, right_idx] = np.random.uniform(
np.clip(right_min_rad, low[right_idx], high[right_idx]),
np.clip(right_max_rad, low[right_idx], high[right_idx]),
size=(n,),
)
def _randomize_symmetric_arms(
self, qpos: np.ndarray, n: int, num_dof_pos: int, low: np.ndarray, high: np.ndarray
) -> None:
# Default range when model joint_limits are missing (low/high are ±inf);
# np.random.uniform requires finite bounds.
default_lo, default_hi = -np.pi, np.pi
for left_idx, right_idx in self._symmetric_arm_pairs:
if left_idx < num_dof_pos and right_idx < num_dof_pos:
lo_l = low[left_idx] if np.isfinite(low[left_idx]) else default_lo
hi_l = high[left_idx] if np.isfinite(high[left_idx]) else default_hi
lo_r = low[right_idx] if np.isfinite(low[right_idx]) else default_lo
hi_r = high[right_idx] if np.isfinite(high[right_idx]) else default_hi
left_range = hi_l - lo_l
left_margin = left_range * self._arm_margin_factor
left_min = lo_l + left_margin
left_max = hi_l - left_margin
right_min = -left_max
right_max = -left_min
right_min_clipped = max(right_min, lo_r)
right_max_clipped = min(right_max, hi_r)
if left_min < left_max:
qpos[:, left_idx] = np.random.uniform(left_min, left_max, size=(n,))
else:
qpos[:, left_idx] = np.random.uniform(lo_l, hi_l, size=(n,))
if right_min_clipped < right_max_clipped:
qpos[:, right_idx] = np.random.uniform(right_min_clipped, right_max_clipped, size=(n,))
else:
qpos[:, right_idx] = np.random.uniform(lo_r, hi_r, size=(n,))
def _randomize_remaining_joints(
self, qpos: np.ndarray, n: int, num_dof_pos: int, low: np.ndarray, high: np.ndarray
) -> None:
used_indices = self._symmetric_arm_used_indices
default_lo, default_hi = -np.pi, np.pi
# 22 = first arm joint (right_shoulder1) in humanoid.xml qpos order;
# arms 22-27 are covered by symmetric_arm_pairs
for i in range(22, num_dof_pos):
if i not in used_indices:
lo = low[i] if np.isfinite(low[i]) else default_lo
hi = high[i] if np.isfinite(high[i]) else default_hi
qpos[:, i] = np.random.uniform(lo, hi, size=(n,))

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

View File

@@ -0,0 +1,95 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os
from dataclasses import dataclass
from motrix_envs import registry
from motrix_envs.base import EnvCfg
bring_ball_model_file = os.path.join(os.path.dirname(__file__), "manipulator_bring_ball.xml")
@registry.envcfg("dm-manipulator-bring-ball")
@dataclass
class BringBallCfg(EnvCfg):
# Simulation
model_file: str = bring_ball_model_file
max_episode_seconds: float = 10.0
sim_dt: float = 0.001
ctrl_dt: float = 0.01
render_spacing: float = 2.5
# Reset sampling (match dm_control defaults).
p_in_hand: float = 0.1
p_in_target: float = 0.1
randomize_arm: bool = True
# Target sampling.
target_x_range: tuple[float, float] = (-0.4, 0.4)
target_z_range: tuple[float, float] = (0.1, 0.4)
target_y: float = 0.001
target_angle_range: tuple[float, float] = (-3.14159265, 3.14159265)
# Object sampling.
object_x_range: tuple[float, float] = (-0.4, 0.4)
object_z_range: tuple[float, float] = (0.0, 0.7)
object_angle_range: tuple[float, float] = (0.0, 6.28318531)
object_x_vel_range: tuple[float, float] = (-5.0, 5.0)
min_object_hand_dist: float = 0.08
# Physics settling at episode start (in control steps, i.e. ctrl_dt units).
# Internally this will be converted to `settle_steps * sim_substeps` physics steps.
settle_steps: int = 80
settle_zero_vel: bool = True
# BringBall reward shaping.
lift_height_threshold: float = 0.04
touch_threshold: float = 0.01
side_penalty_scale: float = 0.05
side_penalty_tanh_scale: float = 10.0
hover_penalty_scale: float = 0.02
hover_close_threshold: float = 0.1
post_grasp_discount: float = 0.7
lift_height_weight: float = 0.3
transport_weight: float = 0.7
transport_progress_scale: float = 0.0
transport_progress_clip: float = 0.02
precision_weight: float = 0.0
precision_margin: float = 0.02
precision_value_at_margin: float = 0.1
# BringBall-specific overrides.
settle_steps: int = 300
p_in_hand: float = 0.0
p_in_target: float = 0.0
randomize_arm: bool = False
object_z_range: tuple[float, float] = (0.2, 0.7)
object_x_vel_range: tuple[float, float] = (0.0, 0.0)
hover_penalty_scale: float = 0.03
post_grasp_discount: float = 0.0
lift_height_weight: float = 0.1
transport_weight: float = 2.0
transport_progress_scale: float = 2.0
transport_progress_clip: float = 0.02
precision_weight: float = 1.0
precision_margin: float = 0.01
# Reward component weights (total reward mixing).
reach_weight: float = 1.0
orient_weight: float = 1.5
pause_weight: float = 0.5
close_weight: float = 2.0
lift_reward_weight: float = 6.0

View File

@@ -0,0 +1,171 @@
<mujoco model="planar manipulator - bring ball">
<include file="../../common/visual.xml"/>
<include file="../../common/skybox.xml"/>
<include file="../../common/materials.xml"/>
<asset>
<texture name="background" type="2d" file="../../common/motphys-ground.png"/>
<material name="background" texture="background" texrepeat="1 1" texuniform="true"/>
</asset>
<visual>
<map shadowclip=".5"/>
<quality shadowsize="2048"/>
</visual>
<option timestep="0.001" cone="elliptic"/>
<default>
<geom friction=".7" solimp="0.9 0.97 0.001" solref=".005 1"/>
<joint solimplimit="0 0.99 0.01" solreflimit=".005 1"/>
<general ctrllimited="true"/>
<tendon width="0.01"/>
<site size=".003 .003 .003" material="site" group="3"/>
<default class="arm">
<geom type="capsule" material="self" density="500"/>
<joint type="hinge" pos="0 0 0" axis="0 -1 0" limited="true"/>
<default class="hand">
<joint damping=".5" range="-10 60"/>
<geom size=".008"/>
<site type="box" size=".018 .005 .005" pos=".022 0 -.002" euler="0 15 0" group="4"/>
<default class="fingertip">
<geom type="sphere" size=".008" material="effector"/>
<joint damping=".01" stiffness=".01" range="-40 20"/>
<site size=".012 .005 .008" pos=".003 0 .003" group="4" euler="0 0 0"/>
</default>
</default>
</default>
<default class="object">
<geom material="self"/>
</default>
<default class="task">
<site rgba="0 0 0 0"/>
</default>
<default class="obstacle">
<geom material="decoration" friction="0"/>
</default>
<default class="ghost">
<geom material="target" contype="0" conaffinity="0"/>
</default>
</default>
<worldbody>
<!-- Arena -->
<light name="light" directional="true" diffuse=".6 .6 .6" pos="0 0 1" dir="0 0.1 -1" specular=".3 .3 .3"/>
<geom name="floor" type="plane" pos="0 0 0" size=".4 .2 10" material="background"/>
<geom name="wall1" type="plane" pos="-.682843 0 .282843" size=".4 .2 10" material="background" zaxis="1 0 1" contype="0" conaffinity="0"/>
<geom name="wall2" type="plane" pos=".682843 0 .282843" size=".4 .2 10" material="background" zaxis="-1 0 1" contype="0" conaffinity="0"/>
<geom name="wall1_collider" type="box" pos="-.7394115 0 .2262745" size="5 .2 .08" quat="0.9238795 0 0.3826834 0" contype="1" conaffinity="1" rgba="0 0 0 0.001"/>
<geom name="wall2_collider" type="box" pos=".7394115 0 .2262745" size="5 .2 .08" quat="0.9238795 0 -0.3826834 0" contype="1" conaffinity="1" rgba="0 0 0 0.001"/>
<geom name="background" type="plane" pos="0 .2 .5" size="1 .5 10" material="background" zaxis="0 -1 0" contype="0" conaffinity="0"/>
<camera name="fixed" pos="0 -16 .4" xyaxes="1 0 0 0 0 1" fovy="4"/>
<!-- Arm -->
<geom name="arm_root" type="cylinder" fromto="0 -.022 .4 0 .022 .4" size=".024"
material="decoration" contype="0" conaffinity="0"/>
<body name="upper_arm" pos="0 0 .4" childclass="arm">
<joint name="arm_root" damping="2" limited="false"/>
<geom name="upper_arm" size=".02" fromto="0 0 0 0 0 .18"/>
<body name="middle_arm" pos="0 0 .18" childclass="arm">
<joint name="arm_shoulder" damping="1.5" range="-160 160"/>
<geom name="middle_arm" size=".017" fromto="0 0 0 0 0 .15"/>
<body name="lower_arm" pos="0 0 .15">
<joint name="arm_elbow" damping="1" range="-160 160"/>
<geom name="lower_arm" size=".014" fromto="0 0 0 0 0 .12"/>
<body name="hand" pos="0 0 .12">
<joint name="arm_wrist" damping=".5" range="-140 140" />
<geom name="hand" size=".011" fromto="0 0 0 0 0 .03"/>
<geom name="palm1" fromto="0 0 .03 .03 0 .045" class="hand"/>
<geom name="palm2" fromto="0 0 .03 -.03 0 .045" class="hand"/>
<site name="grasp" pos="0 0 .065"/>
<body name="pinch site" pos="0 0 .090">
<site name="pinch"/>
<inertial pos="0 0 0" mass="1e-6" diaginertia="1e-12 1e-12 1e-12"/>
<camera name="hand" pos="0 -.3 0" xyaxes="1 0 0 0 0 1" mode="track"/>
</body>
<site name="palm_touch" type="box" group="4" size=".025 .005 .008" pos="0 0 .043"/>
<body name="thumb" pos=".03 0 .045" euler="0 -90 0" childclass="hand">
<joint name="thumb"/>
<geom name="thumb1" fromto="0 0 0 .02 0 -.01" size=".007"/>
<geom name="thumb2" fromto=".02 0 -.01 .04 0 -.01" size=".007"/>
<site name="thumb_touch" group="4"/>
<body name="thumbtip" pos=".05 0 -.01" childclass="fingertip">
<joint name="thumbtip"/>
<geom name="thumbtip1" pos="-.003 0 0" />
<geom name="thumbtip2" pos=".003 0 0" />
<site name="thumbtip_touch" group="4"/>
</body>
</body>
<body name="finger" pos="-.03 0 .045" euler="0 90 180" childclass="hand">
<joint name="finger"/>
<geom name="finger1" fromto="0 0 0 .02 0 -.01" size=".007" />
<geom name="finger2" fromto=".02 0 -.01 .04 0 -.01" size=".007"/>
<site name="finger_touch"/>
<body name="fingertip" pos=".05 0 -.01" childclass="fingertip">
<joint name="fingertip"/>
<geom name="fingertip1" pos="-.003 0 0" />
<geom name="fingertip2" pos=".003 0 0" />
<site name="fingertip_touch"/>
</body>
</body>
</body>
</body>
</body>
</body>
<!-- prop: ball -->
<body name="ball" pos=".4 0 .4" childclass="object">
<joint name="ball_x" type="slide" axis="1 0 0" ref=".4"/>
<joint name="ball_z" type="slide" axis="0 0 1" ref=".4"/>
<joint name="ball_y" type="hinge" axis="0 1 0"/>
<geom name="ball" type="sphere" size=".022" />
<site name="ball" type="sphere"/>
</body>
<!-- target -->
<body name="target_ball" pos=".4 .001 .4" mocap="true" childclass="ghost">
<geom name="target_ball" type="sphere" size=".02" />
<site name="target_ball" type="sphere"/>
</body>
</worldbody>
<tendon>
<fixed name="grasp">
<joint joint="thumb" coef=".5"/>
<joint joint="finger" coef=".5"/>
</fixed>
<fixed name="coupling">
<joint joint="thumb" coef="-.5"/>
<joint joint="finger" coef=".5"/>
</fixed>
</tendon>
<equality>
<tendon name="coupling" tendon1="coupling" solimp="0.95 0.99 0.001" solref=".005 .5"/>
</equality>
<sensor>
<touch name="palm_touch" site="palm_touch"/>
<touch name="finger_touch" site="finger_touch"/>
<touch name="thumb_touch" site="thumb_touch"/>
<touch name="fingertip_touch" site="fingertip_touch"/>
<touch name="thumbtip_touch" site="thumbtip_touch"/>
</sensor>
<actuator>
<motor name="root" joint="arm_root" ctrlrange="-1 1" gear="12"/>
<motor name="shoulder" joint="arm_shoulder" ctrlrange="-1 1" gear="8"/>
<motor name="elbow" joint="arm_elbow" ctrlrange="-1 1" gear="4"/>
<motor name="wrist" joint="arm_wrist" ctrlrange="-1 1" gear="2"/>
<motor name="grasp" tendon="grasp" ctrlrange="-1 1" gear="2"/>
</actuator>
</mujoco>

View File

@@ -0,0 +1,579 @@
# 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.manipulator.cfg import BringBallCfg
from motrix_envs.math import quaternion
from motrix_envs.np import reward as reward_utils
from motrix_envs.np.env import NpEnv, NpEnvState
_ARM_JOINTS = (
"arm_root",
"arm_shoulder",
"arm_elbow",
"arm_wrist",
"finger",
"fingertip",
"thumb",
"thumbtip",
)
_TOUCH_SENSORS = ("palm_touch", "finger_touch", "thumb_touch", "fingertip_touch", "thumbtip_touch")
_HAND_GEOMS = (
"hand",
"palm1",
"palm2",
"thumb1",
"thumb2",
"thumbtip1",
"thumbtip2",
"finger1",
"finger2",
"fingertip1",
"fingertip2",
)
def _sanitize_joint_limits(low: np.ndarray, high: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
low = low.copy()
high = high.copy()
low = np.where(np.isfinite(low), low, -np.pi)
high = np.where(np.isfinite(high), high, np.pi)
return low.astype(np.float32), high.astype(np.float32)
def _quat_from_y_angle(angle: np.ndarray) -> np.ndarray:
zeros = np.zeros_like(angle)
return quaternion.from_euler(zeros, angle, zeros)
def _quat_to_z_axis(quat: np.ndarray) -> np.ndarray:
quat = np.asarray(quat, dtype=np.float32)
return quaternion.rotate_vector(quat, np.array([0.0, 0.0, 1.0], dtype=np.float32)).astype(np.float32)
def _tolerance(
x: np.ndarray,
*,
bounds: tuple[float, float] = (0.0, 0.0),
margin: float = 0.0,
sigmoid: str = "gaussian",
value_at_margin: float = 0.1,
) -> np.ndarray:
"""Vectorized tolerance reward (ported from dm_control-style reward_utils)."""
return reward_utils.tolerance(x, bounds=bounds, margin=margin, sigmoid=sigmoid, value_at_margin=value_at_margin)
class ManipulatorBase(NpEnv):
_cfg: BringBallCfg
_observation_space: gym.spaces.Box
_action_space: gym.spaces.Box
def __init__(self, cfg: BringBallCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
self._cfg = cfg
self._joint_limit_low, self._joint_limit_high = _sanitize_joint_limits(*self._model.joint_limits)
self._arm_joint_pos_indices = np.array([self._joint_pos_index(n) for n in _ARM_JOINTS], dtype=np.int32)
self._arm_joint_vel_indices = np.array([self._joint_vel_index(n) for n in _ARM_JOINTS], dtype=np.int32)
self._thumb_qpos_i = self._joint_pos_index("thumb")
self._finger_qpos_i = self._joint_pos_index("finger")
self._thumbtip_qpos_i = self._joint_pos_index("thumbtip")
self._fingertip_qpos_i = self._joint_pos_index("fingertip")
self._grasp_site = self._model.get_site("grasp")
# Ensure correct actuator index is retrieved from base model
self._grasp_act_i = int(self._model.get_actuator_index("grasp"))
self._object_site = self._model.get_site("ball")
self._target_site = self._model.get_site("target_ball")
target_body = self._model.get_body("target_ball")
if target_body is None:
raise ValueError("Target body 'target_ball' not found in model")
self._target_mocap = target_body.mocap
object_qpos_joints = ("ball_x", "ball_z", "ball_y")
object_geom_names = ("ball",)
self._object_qpos_indices = np.array([self._joint_pos_index(n) for n in object_qpos_joints], dtype=np.int32)
self._object_qvel_indices = np.array([self._joint_vel_index(n) for n in object_qpos_joints], dtype=np.int32)
self._object_x_qvel_i = int(self._object_qvel_indices[0])
self._hand_geom_indices = np.array([self._model.get_geom_index(name) for name in _HAND_GEOMS], dtype=np.uint32)
object_geom_indices = np.array(
[self._model.get_geom_index(name) for name in object_geom_names], dtype=np.uint32
)
self._hand_object_pairs = np.stack(
[
np.repeat(self._hand_geom_indices, object_geom_indices.shape[0]),
np.tile(object_geom_indices, self._hand_geom_indices.shape[0]),
],
axis=-1,
).astype(np.uint32)
self._num_hand_object_pairs = int(self._hand_object_pairs.shape[0])
self._init_dof_pos = self._model.compute_init_dof_pos().astype(np.float32)
self._init_dof_vel = np.zeros((self._model.num_dof_vel,), dtype=np.float32)
self._init_action_space()
self._init_obs_space()
def _joint_pos_index(self, joint_name: str) -> int:
joint_index = self._model.get_joint_index(joint_name)
return int(self._model.joint_dof_pos_indices[joint_index])
def _joint_vel_index(self, joint_name: str) -> int:
joint_index = self._model.get_joint_index(joint_name)
return int(self._model.joint_dof_vel_indices[joint_index])
def _init_action_space(self):
low, high = self._model.actuator_ctrl_limits
self._action_space = gym.spaces.Box(low, high, (self._model.num_actuators,), dtype=np.float32)
def _init_obs_space(self):
# arm_pos(sin,cos)=16 + arm_vel=8 + touch=5 + hand=3 + object=3 + target=3 + rel=3
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (41,), 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: np.ndarray, state: NpEnvState) -> NpEnvState:
actions = np.asarray(actions, dtype=np.float32)
# Enforce actuator control limits to avoid out-of-range impulses.
actions = np.clip(actions, self._action_space.low, self._action_space.high).astype(np.float32)
state.info["last_actions"] = state.info["actions"]
state.info["actions"] = actions
state.data.actuator_ctrls = actions
return state
def _touch_raw(self, data: mtx.SceneData) -> np.ndarray:
values = []
for name in _TOUCH_SENSORS:
v = np.asarray(self._model.get_sensor_value(name, data)).reshape(data.shape[0], -1)[:, 0]
values.append(v)
return np.stack(values, axis=-1).astype(np.float32)
def _touch_log(self, data: mtx.SceneData) -> np.ndarray:
return np.log1p(self._touch_raw(data))
def _hand_pos(self, data: mtx.SceneData) -> np.ndarray:
return self._grasp_site.get_position(data).astype(np.float32)
def _object_pos(self, data: mtx.SceneData) -> np.ndarray:
return self._object_site.get_position(data).astype(np.float32)
def _target_pos(self, data: mtx.SceneData) -> np.ndarray:
return self._target_site.get_position(data).astype(np.float32)
def _contact_with_object(self, data: mtx.SceneData) -> np.ndarray:
cquery = self._model.get_contact_query(data)
colliding = cquery.is_colliding(self._hand_object_pairs)
colliding = np.asarray(colliding).reshape((data.shape[0], self._num_hand_object_pairs))
return colliding.any(axis=-1)
def _get_obs(self, data: mtx.SceneData) -> np.ndarray:
qpos = data.dof_pos[:, self._arm_joint_pos_indices]
arm_pos = np.stack([np.sin(qpos), np.cos(qpos)], axis=-1).reshape(data.shape[0], -1)
arm_vel = data.dof_vel[:, self._arm_joint_vel_indices]
touch = self._touch_log(data)
hand_pos = self._hand_pos(data)
object_pos = self._object_pos(data)
target_pos = self._target_pos(data)
rel = object_pos - target_pos
obs = np.concatenate([arm_pos, arm_vel, touch, hand_pos, object_pos, target_pos, rel], axis=-1)
assert obs.shape == (data.shape[0], self._observation_space.shape[0])
return obs.astype(np.float32)
def _sample_arm_joint_angles(self, num: int) -> np.ndarray:
joint_indices = np.array([self._model.get_joint_index(n) for n in _ARM_JOINTS], dtype=np.int32)
low = self._joint_limit_low[joint_indices]
high = self._joint_limit_high[joint_indices]
return np.random.uniform(low=low, high=high, size=(num, joint_indices.shape[0])).astype(np.float32)
def _sample_target_pose(self, num: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
cfg = self._cfg
target_x = np.random.uniform(cfg.target_x_range[0], cfg.target_x_range[1], size=(num,)).astype(np.float32)
target_z = np.random.uniform(cfg.target_z_range[0], cfg.target_z_range[1], size=(num,)).astype(np.float32)
target_angle = np.random.uniform(cfg.target_angle_range[0], cfg.target_angle_range[1], size=(num,)).astype(
np.float32
)
return target_x, target_z, target_angle
def _set_target_mocap(
self,
data: mtx.SceneData,
target_x: np.ndarray,
target_z: np.ndarray,
target_angle: np.ndarray,
):
pose = np.zeros((data.shape[0], 7), dtype=np.float32)
pose[:, 0] = target_x
pose[:, 1] = float(self._cfg.target_y)
pose[:, 2] = target_z
pose[:, 3:7] = _quat_from_y_angle(target_angle)
self._target_mocap.set_pose(data, pose)
def _set_object_state(
self,
dof_pos: np.ndarray,
dof_vel: np.ndarray,
target_x: np.ndarray,
target_z: np.ndarray,
target_angle: np.ndarray,
grasp_pos: np.ndarray,
):
cfg = self._cfg
num = dof_pos.shape[0]
# Default: uniform in workspace.
object_x = np.random.uniform(cfg.object_x_range[0], cfg.object_x_range[1], size=(num,)).astype(np.float32)
object_z = np.random.uniform(cfg.object_z_range[0], cfg.object_z_range[1], size=(num,)).astype(np.float32)
object_angle = np.random.uniform(cfg.object_angle_range[0], cfg.object_angle_range[1], size=(num,)).astype(
np.float32
)
# dm_control-style object init distribution.
r = np.random.uniform(0.0, 1.0, size=(num,)).astype(np.float32)
in_hand = r < float(cfg.p_in_hand)
in_target = (r >= float(cfg.p_in_hand)) & (r < float(cfg.p_in_hand + cfg.p_in_target))
uniform = ~(in_hand | in_target)
# Avoid initializing the object too close to the hand to prevent interpenetration / impulse explosions.
min_dist = float(getattr(cfg, "min_object_hand_dist", 0.0))
if min_dist > 0.0 and uniform.any():
min_dist_sq = np.float32(min_dist * min_dist)
max_attempts = 50
pending = uniform.copy()
for _ in range(max_attempts):
if not pending.any():
break
dx = object_x - grasp_pos[:, 0]
dz = object_z - grasp_pos[:, 2]
too_close = (dx * dx + dz * dz) < min_dist_sq
pending = pending & too_close
if not pending.any():
break
n = int(pending.sum())
object_x[pending] = np.random.uniform(cfg.object_x_range[0], cfg.object_x_range[1], size=(n,)).astype(
np.float32
)
object_z[pending] = np.random.uniform(cfg.object_z_range[0], cfg.object_z_range[1], size=(n,)).astype(
np.float32
)
object_x[in_target] = target_x[in_target]
object_z[in_target] = target_z[in_target]
object_angle[in_target] = target_angle[in_target]
object_x[in_hand] = grasp_pos[in_hand, 0]
object_z[in_hand] = grasp_pos[in_hand, 2]
object_angle[in_hand] = 0.0
dof_pos[:, self._object_qpos_indices] = np.stack([object_x, object_z, object_angle], axis=-1)
dof_vel[:, self._object_qvel_indices] = 0.0
if uniform.any():
dof_vel[uniform, self._object_x_qvel_i] = np.random.uniform(
cfg.object_x_vel_range[0], cfg.object_x_vel_range[1], size=(int(uniform.sum()),)
).astype(np.float32)
def _settle(self, data: mtx.SceneData):
control_steps = int(self._cfg.settle_steps)
if control_steps <= 0:
return
substeps = int(self._cfg.sim_substeps)
physics_steps = control_steps * max(substeps, 1)
data.actuator_ctrls = np.zeros((data.shape[0], self._model.num_actuators), dtype=np.float32)
for _ in range(physics_steps):
self._model.step(data)
if self._cfg.settle_zero_vel:
data.set_dof_vel(np.zeros((data.shape[0], self._model.num_dof_vel), dtype=np.float32))
self._model.forward_kinematic(data)
def initialize_episode(self, data: mtx.SceneData) -> None:
"""Episode initialization with optional physics settling (dm_control-style)."""
num = int(data.shape[0])
dof_pos = np.tile(self._init_dof_pos, (num, 1))
dof_vel = np.tile(self._init_dof_vel, (num, 1))
# Optionally randomize arm joint angles and symmetrize the hand.
if getattr(self._cfg, "randomize_arm", True):
arm_angles = self._sample_arm_joint_angles(num)
dof_pos[:, self._arm_joint_pos_indices] = arm_angles
dof_pos[:, self._finger_qpos_i] = dof_pos[:, self._thumb_qpos_i]
dof_pos[:, self._fingertip_qpos_i] = dof_pos[:, self._thumbtip_qpos_i]
data.reset(self._model)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
target_x, target_z, target_angle = self._sample_target_pose(num)
self._set_target_mocap(data, target_x, target_z, target_angle)
self._model.forward_kinematic(data)
grasp_pos = self._grasp_site.get_position(data)
self._set_object_state(dof_pos, dof_vel, target_x, target_z, target_angle, grasp_pos)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
arm_qpos = data.dof_pos[:, self._arm_joint_pos_indices].copy()
self._settle(data)
if not getattr(self._cfg, "randomize_arm", True):
dof_pos_after = data.dof_pos.copy()
dof_vel_after = data.dof_vel.copy()
dof_pos_after[:, self._arm_joint_pos_indices] = arm_qpos
dof_vel_after[:, self._arm_joint_vel_indices] = 0.0
data.set_dof_pos(dof_pos_after, self._model)
data.set_dof_vel(dof_vel_after)
self._model.forward_kinematic(data)
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
num = int(data.shape[0])
self.initialize_episode(data)
obs = self._get_obs(data)
info = {
"actions": np.zeros((num, self._model.num_actuators), dtype=np.float32),
"last_actions": np.zeros((num, self._model.num_actuators), dtype=np.float32),
}
return obs, info
@registry.env("dm-manipulator-bring-ball", "np")
class BringBall(ManipulatorBase):
_cfg: BringBallCfg
def __init__(self, cfg: BringBallCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
# 1. Sensors setup
self._fingertip_site = self._model.get_site("fingertip_touch")
self._thumbtip_site = self._model.get_site("thumbtip_touch")
self._touch_idx_palm = _TOUCH_SENSORS.index("palm_touch")
self._touch_idx_fingertip = _TOUCH_SENSORS.index("fingertip_touch")
self._touch_idx_thumbtip = _TOUCH_SENSORS.index("thumbtip_touch")
def _compute_hand_direction(self, data: mtx.SceneData) -> np.ndarray:
"""Calculates the Z-axis vector of the hand (grasp site)."""
grasp_pose = self._grasp_site.get_pose(data)
return _quat_to_z_axis(grasp_pose[:, 3:])
def _get_tip_positions(self, data: mtx.SceneData) -> tuple[np.ndarray, np.ndarray]:
fingertip_pos = self._fingertip_site.get_position(data).astype(np.float32)
thumbtip_pos = self._thumbtip_site.get_position(data).astype(np.float32)
return fingertip_pos, thumbtip_pos
def _compute_aim_direction(self, object_pos: np.ndarray, grasp_pos: np.ndarray) -> np.ndarray:
vec_to_aim = object_pos - grasp_pos
dist_to_aim = np.linalg.norm(vec_to_aim, axis=-1, keepdims=True)
return vec_to_aim / (dist_to_aim + 1e-6)
def _strict_grasp_condition(self, data: mtx.SceneData, object_pos: np.ndarray) -> np.ndarray:
cfg = self._cfg
height_ok = object_pos[:, 2] > float(cfg.lift_height_threshold)
all_touch = self._touch_raw(data)
touch_threshold = float(cfg.touch_threshold)
touch_ok = (
(all_touch[..., self._touch_idx_palm] > touch_threshold)
| (all_touch[..., self._touch_idx_fingertip] > touch_threshold)
| (all_touch[..., self._touch_idx_thumbtip] > touch_threshold)
)
object_contact_ok = self._contact_with_object(data)
return height_ok & touch_ok & object_contact_ok
def update_state(self, state: NpEnvState) -> NpEnvState:
data = state.data
cfg = self._cfg
# 1. Observation
obs = self._get_obs(data)
terminated = np.isnan(obs).any(axis=-1)
# 2. Positions
object_pos = self._object_pos(data)
target_pos = self._target_pos(data)
grasp_pos = self._hand_pos(data)
# 3. Kinematics
fingertip_pos, thumbtip_pos = self._get_tip_positions(data)
dist_finger = np.linalg.norm(fingertip_pos - object_pos, axis=-1)
dist_thumb = np.linalg.norm(thumbtip_pos - object_pos, axis=-1)
avg_tip_dist = ((dist_finger + dist_thumb) / 2.0).astype(np.float32)
move_dist = np.linalg.norm(object_pos - target_pos, axis=-1).astype(np.float32)
# 4. Dynamics
arm_vel = data.dof_vel[:, self._arm_joint_vel_indices[:4]].astype(np.float32)
arm_speed = np.linalg.norm(arm_vel, axis=-1).astype(np.float32)
arm_speed_step = (arm_speed * float(cfg.ctrl_dt)).astype(np.float32)
# 5. Logic Checks
is_grasped = self._strict_grasp_condition(data, object_pos)
contact_with_obj = self._contact_with_object(data)
hover_threshold = float(cfg.hover_close_threshold)
is_close_to_ball = (avg_tip_dist < hover_threshold).astype(np.float32)
grasp_mask = is_grasped.astype(np.float32)
post_grasp_scale = 1.0 - grasp_mask * float(cfg.post_grasp_discount)
# --- Rewards ---
# R1: Reach
r_reach = _tolerance(avg_tip_dist, bounds=(0.0, 0.02), margin=0.25, sigmoid="linear").astype(np.float32)
r_reach = (r_reach * post_grasp_scale).astype(np.float32)
# R2: Orient
hand_dir = self._compute_hand_direction(data)
unit_vec_to_aim = self._compute_aim_direction(object_pos, grasp_pos)
pointing_dot = np.sum(hand_dir * unit_vec_to_aim, axis=-1)
# Dynamic tolerance
dist_from_base = np.linalg.norm(object_pos[:, :2], axis=-1)
orient_bound_lower = 0.95 * np.clip(dist_from_base / 0.5, 0.0, 1.0)
r_orient_raw = 1.0 - orient_bound_lower + pointing_dot
r_orient = np.clip(r_orient_raw, 0.0, 1.0).astype(np.float32)
r_orient = (r_orient * post_grasp_scale).astype(np.float32)
# R3: Pause
r_pause = (
_tolerance(arm_speed_step, bounds=(0.0, 0.05), margin=0.3, sigmoid="linear").astype(np.float32)
* is_close_to_ball
)
r_pause = (r_pause * post_grasp_scale).astype(np.float32)
# R4: Close
default_actions = np.zeros((data.shape[0], self._model.num_actuators), dtype=np.float32)
grasp_action = state.info.get("actions", default_actions)[:, self._grasp_act_i].astype(np.float32)
r_close_intent = _tolerance(
grasp_action, bounds=(0.8, 1.0), margin=1.0, sigmoid="linear", value_at_margin=0.01
).astype(np.float32)
r_approach_grasp = r_close_intent * is_close_to_ball * r_orient * r_pause * contact_with_obj.astype(np.float32)
r_sustain_grasp = r_close_intent * grasp_mask
r_close = (r_approach_grasp * (1.0 - grasp_mask) + r_sustain_grasp).astype(np.float32)
# R5: Lift & Transport
lift_h = float(cfg.lift_height_threshold)
ball_z = object_pos[:, 2].astype(np.float32)
r_lift_height = (
_tolerance(ball_z, bounds=(lift_h, lift_h + 0.15), margin=0.02, sigmoid="linear", value_at_margin=0.01)
* grasp_mask
).astype(np.float32)
r_transport = (_tolerance(move_dist, bounds=(0.0, 0.01), margin=0.3, sigmoid="linear") * grasp_mask).astype(
np.float32
)
r_precision = (
_tolerance(
move_dist,
bounds=(0.0, 0.0),
margin=float(cfg.precision_margin),
sigmoid="gaussian",
value_at_margin=float(cfg.precision_value_at_margin),
)
* grasp_mask
).astype(np.float32)
lift_height_weight = float(cfg.lift_height_weight)
transport_weight = float(cfg.transport_weight)
lift_norm = max(lift_height_weight + transport_weight, 1e-6)
r_lift = ((lift_height_weight * r_lift_height + transport_weight * r_transport) / lift_norm).astype(np.float32)
prev_move_dist = state.info.get("prev_move_dist")
if prev_move_dist is None:
prev_move_dist = move_dist
else:
prev_move_dist = np.asarray(prev_move_dist, dtype=np.float32)
if "steps" in state.info:
first_step = state.info["steps"] == 0
prev_move_dist = np.where(first_step, move_dist, prev_move_dist)
progress_clip = float(cfg.transport_progress_clip)
progress = (prev_move_dist - move_dist) / max(progress_clip, 1e-6)
progress = np.clip(progress, -1.0, 1.0).astype(np.float32)
r_progress = (progress * float(cfg.transport_progress_scale) * grasp_mask).astype(np.float32)
state.info["prev_move_dist"] = move_dist.astype(np.float32)
# --- Penalties ---
all_touch = self._touch_raw(data)
side_touch_sum = (all_touch[..., 1] + all_touch[..., 2]).astype(np.float32)
penalty_side = (
-float(cfg.side_penalty_scale) * np.tanh(side_touch_sum * float(cfg.side_penalty_tanh_scale))
).astype(np.float32)
hover_phase = (is_close_to_ball > 0.5) & (~contact_with_obj)
penalty_hover = (-float(cfg.hover_penalty_scale) * hover_phase.astype(np.float32)).astype(np.float32)
# --- Total ---
reach_w = float(cfg.reach_weight)
orient_w = float(cfg.orient_weight)
pause_w = float(cfg.pause_weight)
close_w = float(cfg.close_weight)
lift_w = float(cfg.lift_reward_weight)
precision_w = float(cfg.precision_weight)
weight_sum = max(reach_w + orient_w + pause_w + close_w + lift_w + precision_w, 1e-6)
reward = (
(
reach_w * r_reach
+ orient_w * r_orient
+ pause_w * r_pause
+ close_w * r_close
+ lift_w * r_lift
+ precision_w * r_precision
)
/ weight_sum
+ penalty_side
+ penalty_hover
+ r_progress
).astype(np.float32)
reward = np.where(terminated, 0.0, reward)
state.info["Reward"] = {
"reach": r_reach,
"orient": r_orient,
"close": r_close,
"lift": r_lift,
"transport": r_transport,
"precision": r_precision,
"progress": r_progress,
"total": reward,
}
state.info["metrics"] = {
"pointing_dot": pointing_dot,
"is_grasped": is_grasped.astype(np.float32),
"avg_tip_dist": avg_tip_dist,
"move_dist": move_dist,
"transport_reward": r_transport,
"precision_reward": r_precision,
"progress_reward": r_progress,
}
return state.replace(obs=obs, reward=reward, terminated=terminated)

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

View File

@@ -0,0 +1,39 @@
# 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
import numpy as np
from motrix_envs import registry
from motrix_envs.base import EnvCfg
model_file = os.path.dirname(__file__) + "/pendulum.xml"
@registry.envcfg("pendulum")
@dataclass
class PendulumEnvCfg(EnvCfg):
model_file: str = model_file
max_episode_seconds: float = 10.0
sim_dt: float = 0.0125
ctrl_dt: float = 0.025
angle_bound: float = 8.0
cosing_bound: float = 0.0
# reset_noise_scale: float = 0.01
def __post_init__(self):
self.cosing_bound = float(np.cos(np.deg2rad(self.angle_bound)))

View File

@@ -0,0 +1,35 @@
<mujoco model="pendulum">
<include file="../../common/visual.xml" />
<include file="../../common/skybox.xml" />
<include file="../../common/materials.xml" />
<asset>
<texture name="motphys_ground_tex" type="2d" file="../../common/motphys-ground.png" />
<material name="motphys_ground_mat" texture="motphys_ground_tex" texrepeat="1 1" texuniform="true" />
</asset>
<option timestep="0.02">
<flag contact="disable" energy="enable"/>
</option>
<worldbody>
<light diffuse="0.6 0.6 0.6" pos="0 0 1.5"
dir="-0.49835488200187683 0.2925136387348175 -0.8161361217498779" directional="true" />
<geom name="floor" size="0 0 0.2" type="plane" material="motphys_ground_mat"/>
<camera name="fixed" pos="0 -1.5 2" xyaxes='1 0 0 0 1 1'/>
<camera name="lookat" mode="targetbodycom" target="pole" pos="0 -2 1"/>
<body name="pole" pos="0 0 .6">
<joint name="hinge" type="hinge" axis="0 1 0" damping="0.1"/>
<geom name="base" material="decoration" type="cylinder" fromto="0 -.03 0 0 .03 0" size="0.021" mass="0"/>
<geom name="pole" material="self" type="capsule" fromto="0 0 0 0 0 0.5" size="0.02" mass="0"/>
<geom name="mass" material="effector" type="sphere" pos="0 0 0.5" size="0.05" mass="1"/>
</body>
</worldbody>
<actuator>
<motor name="torque" joint="hinge" gear="5" ctrlrange="-1 1" ctrllimited="true"/>
</actuator>
</mujoco>

View File

@@ -0,0 +1,117 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.np import reward as reward_utils
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import PendulumEnvCfg
@registry.env("pendulum", "np")
class PendulumEnv(NpEnv):
_cfg: PendulumEnvCfg
def __init__(self, cfg: PendulumEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
ctrl_limits = self._model.actuator_ctrl_limits
self._action_low = float(ctrl_limits[0, 0])
self._action_high = float(ctrl_limits[1, 0])
self._action_space = gym.spaces.Box(-1.0, 1.0, (1,), dtype=np.float32)
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (3,), 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):
actions = np.clip(actions, -1.0, 1.0)
scaled = self._action_low + (actions + 1.0) * 0.5 * (self._action_high - self._action_low)
state.data.actuator_ctrls = scaled
return state
def update_state(self, state: NpEnvState):
# compute observation
data = state.data
dof_pos = data.dof_pos
dof_vel = data.dof_vel
angle = dof_pos[:, 0]
ang_vel = dof_vel[:, 0]
obs = np.stack([np.cos(angle), np.sin(angle), ang_vel], axis=-1)
assert obs.shape == (self._num_envs, 3)
# compute reward
angle_wrapped = (angle + np.pi) % (2 * np.pi) - np.pi
ctrl = data.actuator_ctrls[:, 0]
# In this model, zero angle corresponds to the hanging-down position.
# Shift the target by pi to encourage the upright (inverted) posture.
upright = (1.0 + np.cos(angle_wrapped)) * 0.5
prev_ctrl = state.info.get("prev_ctrl", np.zeros_like(ctrl))
ctrl_delta = ctrl - prev_ctrl
vel_penalty = 0.2 * (ang_vel**2)
energy = 0.5 * ang_vel**2 + (1.0 - np.cos(angle_wrapped))
energy_target = 2.0
energy_reward = reward_utils.tolerance(
energy,
bounds=(energy_target, energy_target),
margin=2.0,
value_at_margin=0.1,
sigmoid="gaussian",
)
reward = (3.0 * upright + energy_reward - vel_penalty - 0.001 * ctrl**2 - 0.001 * ctrl_delta**2).astype(
np.float32
)
# compute terminated
terminated = np.isnan(obs).any(axis=-1)
state.obs = obs
state.reward = reward
state.terminated = terminated
state.info["prev_ctrl"] = ctrl
return state
def reset(self, data: mtx.SceneData):
cfg: PendulumEnvCfg = self._cfg
reset_noise_scale = getattr(cfg, "reset_noise_scale", 0.0)
num_reset = data.shape[0]
dof_pos = np.zeros((num_reset, self._num_dof_pos), dtype=np.float32)
dof_vel = np.zeros((num_reset, self._num_dof_vel), dtype=np.float32)
dof_pos[:, 0] = np.random.uniform(-np.pi, np.pi, size=(num_reset,))
if reset_noise_scale > 0.0:
dof_vel[:, 0] = np.random.uniform(-reset_noise_scale, reset_noise_scale, size=(num_reset,))
data.reset(self._model)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
angle = dof_pos[:, 0]
ang_vel = dof_vel[:, 0]
obs = np.stack([np.cos(angle), np.sin(angle), ang_vel], axis=-1)
return obs, {"prev_ctrl": np.zeros((num_reset,), dtype=np.float32)}

View File

@@ -19,7 +19,7 @@ import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.math.quaternion import Quaternion
from motrix_envs.math import quaternion
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import AnymalCEnvCfg
@@ -185,7 +185,7 @@ class AnymalCEnv(NpEnv):
# Get commands - convert to relative velocity commands
pose_commands = state.info["pose_commands"]
robot_position = root_pos[:, :2]
robot_heading = Quaternion.get_yaw(root_quat)
robot_heading = quaternion.get_yaw(root_quat)
target_position = pose_commands[:, :2]
target_heading = pose_commands[:, 2]
@@ -295,14 +295,14 @@ class AnymalCEnv(NpEnv):
)
robot_arrow_pos = robot_pos.copy()
robot_arrow_pos[:, 2] = arrow_height
robot_arrow_quat = Quaternion.from_euler(0, 0, cur_yaw)
robot_arrow_quat = quaternion.from_euler(0, 0, cur_yaw)
mocap = self._model.get_body("robot_heading_arrow").mocap
mocap.set_pose(data, np.concatenate([robot_arrow_pos, robot_arrow_quat], axis=1))
des_yaw = np.where(
np.linalg.norm(desired_vel_xy, axis=1) > 1e-6, np.arctan2(desired_vel_xy[:, 1], desired_vel_xy[:, 0]), 0.0
)
desired_arrow_quat = Quaternion.from_euler(0, 0, des_yaw)
desired_arrow_quat = quaternion.from_euler(0, 0, des_yaw)
mocap = self._model.get_body("desired_heading_arrow").mocap
mocap.set_pose(data, np.concatenate([robot_arrow_pos, desired_arrow_quat], axis=1))
@@ -350,7 +350,7 @@ class AnymalCEnv(NpEnv):
# Get robot position and heading for arrival determination
robot_position = pose[:, :2]
robot_heading = Quaternion.get_yaw(root_quat)
robot_heading = quaternion.get_yaw(root_quat)
target_position = info["pose_commands"][:, :2]
target_heading = info["pose_commands"][:, 2]
position_error = target_position - robot_position
@@ -454,7 +454,7 @@ class AnymalCEnv(NpEnv):
arrow_pos = pose_commands.copy()
arrow_pos[:, 2] = 0.05
arrow_pos = np.column_stack([pose_commands[:, 0], pose_commands[:, 1], np.full((num_envs, 1), 0.5)])
arrow_quat = Quaternion.from_euler(0, 0, pose_commands[:, 2])
arrow_quat = quaternion.from_euler(0, 0, pose_commands[:, 2])
mocap = self._model.get_body("target_marker").mocap
mocap.set_pose(data, np.concatenate([arrow_pos, arrow_quat], axis=1))
@@ -490,7 +490,7 @@ class AnymalCEnv(NpEnv):
return state.replace(terminated=terminated)
def reset(self, data: mtx.SceneData, done: np.ndarray = None) -> tuple[np.ndarray, dict]:
def reset(self, data: mtx.SceneData) -> tuple[np.ndarray, dict]:
cfg: AnymalCEnvCfg = self._cfg
num_envs = data.shape[0]
@@ -563,7 +563,7 @@ class AnymalCEnv(NpEnv):
# Calculate velocity commands (consistent with update_state)
robot_position = root_pos[:, :2]
robot_heading = Quaternion.get_yaw(root_quat)
robot_heading = quaternion.get_yaw(root_quat)
target_position = pose_commands[:, :2]
target_heading = pose_commands[:, 2]
@@ -651,4 +651,4 @@ class AnymalCEnv(NpEnv):
def _compute_projected_gravity(self, quat: np.ndarray) -> np.ndarray:
gravity = np.array([0.0, 0.0, -1.0], dtype=np.float32)
return Quaternion.rotate_vector(quat, gravity)
return quaternion.rotate_vector(quat, gravity)

View File

@@ -19,7 +19,7 @@ import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpEnvCfg
from motrix_envs.math.quaternion import Quaternion
from motrix_envs.math import quaternion
from motrix_envs.np.env import NpEnv, NpEnvState
@@ -190,7 +190,7 @@ class Go1WalkTask(NpEnv):
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
local_gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
local_gravity = quaternion.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
@@ -320,7 +320,7 @@ class Go1WalkTask(NpEnv):
# Penalize non flat base orientation
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):

View File

@@ -19,7 +19,7 @@ import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpRoughEnvCfg
from motrix_envs.math.quaternion import Quaternion
from motrix_envs.math import quaternion
from motrix_envs.np.env import NpEnv, NpEnvState
from .common import generate_repeating_array
@@ -44,18 +44,15 @@ class Go1WalkRoughTask(NpEnv):
(self._num_dof_vel,),
dtype=np.float32,
)
self.height_list = np.array([-2.5, 0.5, 2.0])
offset_h = [[2, 2, 1, 1, 1], [2, 2, 1, 1, 2], [2, 1, 1, 1, 1], [2, 1, 1, 1, 1], [2, 1, 1, 1, 1]]
offset = []
for i in range(5):
for j in range(5):
h_index = offset_h[j][i]
offset.append([(i - 2) * 8.0, (j - 2) * 8.0, self.height_list[h_index]])
self.offset_list = np.array(offset)
go1_init_height = 0.3
self.reset_offset = self._generate_init_offsets(go1_init_height)
self._init_dof_pos = self._model.compute_init_dof_pos()
self._init_dof_pos[2] = self.height_list[0]
geom_floor_pos = self._model.get_geom("floor").local_pose[:3]
geom_floor_pos[2] += go1_init_height
self._init_dof_pos[:3] = geom_floor_pos
self._init_buffer()
self.height_counter = 0
self.reset_counter = 0
def _init_obs_space(self):
model = self.model
@@ -216,7 +213,7 @@ class Go1WalkRoughTask(NpEnv):
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
local_gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
local_gravity = quaternion.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
@@ -309,9 +306,9 @@ class Go1WalkRoughTask(NpEnv):
if self.training_level == 1:
num_period = 25
idx = generate_repeating_array(num_period, num_reset, self.height_counter)
self.height_counter = (self.height_counter + num_reset) % num_period
dof_pos[:, :3] = self.offset_list[idx]
idx = generate_repeating_array(num_period, num_reset, self.reset_counter)
self.reset_counter = (self.reset_counter + num_reset) % num_period
dof_pos[:, :3] = self.reset_offset[idx]
data.reset(self._model)
data.set_dof_vel(dof_vel)
@@ -364,7 +361,7 @@ class Go1WalkRoughTask(NpEnv):
# Penalize non flat base orientation
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):
@@ -433,5 +430,48 @@ class Go1WalkRoughTask(NpEnv):
# check whether the robot reaching into the terrain border and change the move direction
border_size = 19.0
position = self._body.get_position(data)
is_out = (np.square(position[:, :2]) > border_size**2).any(axis=1)
geom_floor = self._model.get_geom("floor")
geom_floor_rough = self._model.get_geom("floor_rough")
in_rough_aera = self._is_in_area(position, border_size, border_size, geom_floor_rough.local_pose[:2])
in_flat_aera = self._is_in_area(position, border_size, border_size, geom_floor.local_pose[:2])
is_out = ~(in_rough_aera | in_flat_aera)
info["commands"][is_out] = [0, 0, 0]
def _is_in_area(self, pos, length, width, offset):
x = pos[:, 0]
y = pos[:, 1]
return (
((offset[0] - length) < x)
& ((offset[0] + length) > x)
& ((offset[1] - width) < y)
& ((offset[1] + width) > y)
)
def _generate_init_offsets(self, init_height: float, grid_size: tuple[int, int] = (5, 5)) -> np.ndarray:
"""Generate initialization offsets for rough terrain training."""
hfield = self._model.get_geom("floor_rough").hfield
nx, ny = grid_size
# Vectorized grid generation
idx_x, idx_y = np.meshgrid(np.arange(nx), np.arange(ny, dtype=np.float32), indexing="ij")
idx_x, idx_y = idx_x.flatten(), idx_y.flatten()
# Map to heightfield indices and get heights
hfield_idx_x = (hfield.ncol * idx_x // nx).astype(int)
hfield_idx_y = (hfield.nrow * idx_y // ny).astype(int)
heights = np.array([hfield.get(iy, ix) for ix, iy in zip(hfield_idx_x, hfield_idx_y)])
# Calculate positions
grid_len_x = (hfield.bound[3] - hfield.bound[0]) / nx
grid_len_y = (hfield.bound[4] - hfield.bound[1]) / ny
center_offset = np.array([grid_len_x / 2, grid_len_y / 2])
offsets = np.column_stack(
[
hfield.bound[0] + grid_len_x * idx_x + center_offset[0],
hfield.bound[4] - grid_len_y * idx_y - center_offset[1],
heights + init_height,
]
).astype(np.float32)
return offsets

View File

@@ -19,7 +19,7 @@ import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpStairsEnvCfg
from motrix_envs.math.quaternion import Quaternion
from motrix_envs.math import quaternion
from motrix_envs.np.env import NpEnv, NpEnvState
from .common import generate_repeating_array
@@ -215,7 +215,7 @@ class Go1WalkStairsTask(NpEnv):
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
local_gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
local_gravity = quaternion.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
@@ -277,7 +277,7 @@ class Go1WalkStairsTask(NpEnv):
force = []
for foot in self.cfg.sensor.feet:
contact_force = self._model.get_sensor_value(foot + "_foot_contact", data)
contact_force = Quaternion.rotate_inverse(base_quat, contact_force)
contact_force = quaternion.rotate_inverse(base_quat, contact_force)
force.append(contact_force)
return np.concatenate(force, axis=1)
@@ -370,7 +370,7 @@ class Go1WalkStairsTask(NpEnv):
# Penalize non flat base orientation
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = Quaternion.rotate_inverse(base_quat, self.gravity_vec)
gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):

View File

@@ -14,13 +14,14 @@
<asset>
<hfield name="hfield" file="assets/heightmap.png" size="20 20 2.5 0.1" />
<hfield name="hfield-plane" nrow="2" ncol="2" elevation="0 0 0 0" size="20 20 1.0 0.1" />
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true" />
<geom name="floor" pos="0 0 -3" size="0 0 0.01" type="plane" material="motphys-ground"
contype="1" conaffinity="0" priority="1" friction="0.6" condim="3" />
<geom name="floor_rough" pos="0 0 -1" type="hfield" hfield="hfield" material="motphys-ground"
<geom name="floor_rough" pos="0 0 0" type="hfield" hfield="hfield" material="motphys-ground"
contype="1" conaffinity="0" priority="1" friction="0.6" />
<geom name="floor" pos="0 40 0" type="hfield" hfield="hfield-plane" material="motphys-ground"
contype="1" conaffinity="0" priority="1" friction="0.6" />
</worldbody>
</mujoco>

View File

@@ -41,22 +41,22 @@ class ControlConfig:
@dataclass
class InitState:
# the initial position of the robot in the world frame
pos = [0.0, 0.0, 0.42] #0.278
pos = [0.0, 0.0, 0.42] # 0.278
# the default angles for all joints. key = joint name, value = target angle [rad]
default_joint_angles = {
"FL_hip" : 0.1, # [rad]
"FL_thigh" : 0.9, # [rad]
"FL_calf" : -1.8, # [rad]
"FR_hip" : -0.1, # [rad]
"FR_thigh" : 0.9, # [rad]
"FR_calf" : -1.8, # [rad]
"RL_hip" : 0.1, # [rad]
"RL_thigh" : 0.9, # [rad]
"RL_calf" : -1.8, # [rad]
"RR_hip" : -0.1, # [rad]
"RR_thigh" : 0.9, # [rad]
"RR_calf" : -1.8, # [rad]
"FL_hip": 0.1, # [rad]
"FL_thigh": 0.9, # [rad]
"FL_calf": -1.8, # [rad]
"FR_hip": -0.1, # [rad]
"FR_thigh": 0.9, # [rad]
"FR_calf": -1.8, # [rad]
"RL_hip": 0.1, # [rad]
"RL_thigh": 0.9, # [rad]
"RL_calf": -1.8, # [rad]
"RR_hip": -0.1, # [rad]
"RR_thigh": 0.9, # [rad]
"RR_calf": -1.8, # [rad]
}
@@ -64,7 +64,7 @@ class InitState:
class Commands:
vel_limit = [
[-2.0, -1.0, -3.1416], # min: vel_x [m/s], vel_y [m/s], ang_vel [rad/s]
[ 2.0, 1.0, 3.1416], # max
[2.0, 1.0, 3.1416], # max
]
@@ -82,8 +82,13 @@ class Asset:
foot_name = "foot"
penalize_contacts_on = ["thigh", "calf"]
terminate_after_contacts_on = [
"base_collision_0", "base_collision_1", "base_collision_2",
"fl_hip_0", "fr_hip_0", "rl_hip_0", "rr_hip_0",
"base_collision_0",
"base_collision_1",
"base_collision_2",
"fl_hip_0",
"fr_hip_0",
"rl_hip_0",
"rr_hip_0",
]
ground = "floor"

View File

@@ -19,41 +19,10 @@ import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go2.cfg import Go2WalkNpEnvCfg
from motrix_envs.math import quaternion
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("go2-flat-terrain-walk", sim_backend="np")
class Go2WalkTask(NpEnv):
_init_dof_pos: np.ndarray
@@ -141,9 +110,7 @@ class Go2WalkTask(NpEnv):
self.hip_indices.append(i)
if "calf" in self._model.actuator_names[i]:
self.calf_indices.append(i)
print("Default joint angles:", self.default_angles)
print("Actuator names:", self._model.actuator_names)
self._init_dof_pos[-self._num_action :] = self.default_angles
self.ground = self._model.get_geom_index(cfg.asset.ground)
@@ -213,7 +180,7 @@ class Go2WalkTask(NpEnv):
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)
local_gravity = quaternion.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
@@ -343,7 +310,7 @@ class Go2WalkTask(NpEnv):
# 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)
gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):

View File

@@ -13,4 +13,4 @@
# limitations under the License.
# ==============================================================================
from . import franka_lift_cube, franka_open_cabinet # noqa: F401 import to register envs
from . import franka_lift_cube, franka_open_cabinet, shadow_hand # noqa: F401 import to register envs

View File

@@ -240,11 +240,11 @@ class FrankaLiftCubeEnv(NpEnv):
## action penalty rate
reach_weight = 1.5 # Cannot be too small
cmd_tracking_weight = 10
cmd_tracking_fine_graind_weight = 20 # Should be larger, need strong pull to target area
object_command_tracking_close_reward_weight = 10
cmd_tracking_weight = 10.0
cmd_tracking_fine_graind_weight = 20.0 # Should be larger, need strong pull to target area
object_command_tracking_close_reward_weight = 10.0
if self.count < 10000:
if self.count < 20000:
action_penalty_rate = 1e-4
joint_vel_penalty_rate = 1e-4
else:

View File

@@ -18,7 +18,7 @@ import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.math.quaternion import Quaternion
from motrix_envs.math import quaternion
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import FrankaOpenCabinetEnvCfg
@@ -206,7 +206,7 @@ class FrankaOpenCabinetEnv(NpEnv):
dist_reward *= 10
## matching orientation reward
quat_reward = Quaternion.similarity(robot_grasp_pose[:, -4:], drawer_grasp_pose[:, -4:])
quat_reward = quaternion.similarity(robot_grasp_pose[:, -4:], drawer_grasp_pose[:, -4:])
## close gripper reward
# When gripper distance < 0.025, closing gripper gets reward
@@ -227,6 +227,7 @@ class FrankaOpenCabinetEnv(NpEnv):
open_reward = (
np.bitwise_not(wrong_open) * open_reward
) # No reward for forced opening (can't force open after increasing MJCF resistance)
quat_reward = np.where(open_reward > 0, 1.0, quat_reward)
##################### Penalty Terms #####################"
## Action penalty

View File

@@ -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.
# ==============================================================================
from .cfg import ShadowHandReposeEnvCfg
# Use the full implementation from env.py (157-dim obs with fingertips)
from .shadow_hand_np import ShadowHandReposeEnv
__all__ = ["ShadowHandReposeEnvCfg", "ShadowHandReposeEnv"]

View File

@@ -0,0 +1,142 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Configuration for Shadow Hand Cube Reorientation Environment"""
import os
from dataclasses import dataclass
from typing import List, Tuple
from motrix_envs import registry
from motrix_envs.base import EnvCfg
# Path to the repose_cube.xml model
model_file = os.path.join(os.path.dirname(__file__), "xmls", "repose_cube.xml")
@registry.envcfg("shadow-hand-repose")
@dataclass
class ShadowHandReposeEnvCfg(EnvCfg):
"""
Configuration for Shadow Hand Cube Reorientation Environment
This environment trains the Shadow Hand to reorient a cube to match
randomly sampled target orientations.
"""
# ====================
# Model Configuration
# ====================
model_file: str = model_file
# ====================
# Simulation Parameters
# ====================
sim_dt: float = 0.01
sim_substeps: int = 1 # Number of simulation steps per control step
ctrl_dt: float = sim_dt * sim_substeps
max_episode_seconds: float = 10.0
max_episode_steps: int = int(max_episode_seconds / ctrl_dt)
# ====================
# Robot Configuration
# ====================
num_hand_dofs: int = 24 # Total DOFs in Shadow Hand
num_actuators: int = 20 # Actuated joints
# Fingertip link names for forward kinematics
fingertip_link_names: List[str] = (
"rh_ffdistal", # First finger (index) distal
"rh_mfdistal", # Middle finger distal
"rh_rfdistal", # Ring finger distal
"rh_lfdistal", # Little finger distal
"rh_thdistal", # Thumb distal
)
# ====================
# Object Configuration
# ====================
cube_initial_pos: Tuple[float, float, float] = (0.33, 0.00, 0.295) # Initial cube position
# ====================
# Reward Parameters
# ====================
# Core reward components
dist_reward_scale: float = -10.0 # Balanced for MotrixSim (推荐)
rot_reward_scale: float = 1.0 # Moderate rotation reward
rot_eps: float = 0.1 # Stable denominator
action_penalty_scale: float = -0.0002
# Success and failure criteria
success_tolerance: float = 0.1 # ~8.6° (moderate challenge)
reach_goal_bonus: float = 2.0 # Balanced incentive
fall_dist: float = 0.24 # Reasonable manipulation space # Distance threshold for dropping cube (meters)
fall_penalty: float = 0.0 # Penalty for dropping the cube
# In-hand distance threshold (only used for success check, not reward)
in_hand_dist_threshold: float = 0.05 # Distance threshold for "in-hand" (5cm)
# Success hold mechanism (uses max_consecutive_successes)
max_consecutive_successes: int = 50 # Reset after holding success for this many steps
# Averaging factor for consecutive successes tracking
av_factor: float = 0.1
# ====================
# Reset Noise Parameters
# ====================
reset_position_noise: float = 0.01 # Increased robustness
reset_dof_pos_noise: float = 0.2 # Higher generalization
reset_dof_vel_noise: float = 0.0 # DOF velocity noise at reset
# ====================
# Observation Scaling
# ====================
vel_obs_scale: float = 0.2 # Scale factor for velocity observations
# ====================
# Action Processing
# ====================
act_moving_average: float = 1.0 # Action smoothing (1.0 = no smoothing)
# ====================
# Visualization
# ====================
# Offset for target visualization (relative to hand position)
# Recommended: offset to upper-left to avoid occluding the real hand and cube
viz_target_offset: Tuple[float, float, float] = (
0.0, # Left (negative X)
0.0, # Forward/Up (negative Y)
0.2, # Up (positive Z)
)
# ====================
# Domain Randomization (Optional)
# ====================
# Enable domain randomization (recommended to start with False)
enable_domain_randomization: bool = False
# Randomization parameters (only used if enable_domain_randomization=True)
randomize_friction: bool = False
friction_range: Tuple[float, float] = (0.8, 12)
randomize_mass: bool = False
mass_range: Tuple[float, float] = (0.8, 1.2)
randomize_com: bool = False
com_displacement_range: float = 0.01

View File

@@ -0,0 +1,391 @@
# 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.
# ==============================================================================
"""
Shadow Hand Cube Reorientation Environment for MotrixSim
This environment implements the classic in-hand cube manipulation task where the
Shadow Hand must reorient a cube to match random target orientations.
"""
import gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.math import quaternion, utils
from motrix_envs.np.env import NpEnv, NpEnvState
from .cfg import ShadowHandReposeEnvCfg
@registry.env("shadow-hand-repose", sim_backend="np")
class ShadowHandReposeEnv(NpEnv):
"""
Shadow Hand Cube Reorientation Environment
Observation space: 157 dimensions
- 24: hand dof positions (unscaled)
- 24: hand dof velocities (scaled by 0.2)
- 7: object pose (pos + quat)
- 3: object linear velocity
- 3: object angular velocity (scaled by 0.2)
- 7: goal pose (pos + quat)
- 4: relative quaternion (object to goal)
- 65: fingertip states (5 fingertips * 13: pos + quat + vel)
- 20: previous actions
Action space: 20 dimensions (normalized [-1, 1] position targets for actuators)
"""
_cfg: ShadowHandReposeEnvCfg
def __init__(self, cfg: ShadowHandReposeEnvCfg, num_envs: int = 1):
super().__init__(cfg, num_envs=num_envs)
# Get model info
self._num_hand_dofs = cfg.num_hand_dofs # 24 total DOFs
self._num_actuators = cfg.num_actuators # 20 actuated joints
# Initialize spaces
self._action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(self._num_actuators,), dtype=np.float32)
self._observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(157,), dtype=np.float32)
# Get actuator control ranges from model
self._actuator_ctrl_lower = self._model.actuator_ctrl_limits[0, :]
self._actuator_ctrl_upper = self._model.actuator_ctrl_limits[1, :]
# Get joint limits for all DOFs (use model's joint_limits directly)
self._hand_dof_lower_limits = self._model.joint_limits[0, :]
self._hand_dof_upper_limits = self._model.joint_limits[1, :]
# Fingertip link indices (use get_link_index for pose/velocity access)
self._fingertip_link_ids = []
for name in cfg.fingertip_link_names:
link_id = self._model.get_link_index(name)
self._fingertip_link_ids.append(link_id)
self._num_fingertips = len(self._fingertip_link_ids)
# Get cube and target link indices
self._cube_link_id = self._model.get_link_index("cube")
self._cube_body = self._model.get_body("cube")
self._cube_dof_vel_indices = self._cube_body.get_dof_vel_indices()
# Target is a mocap body - access via Body object (no get_mocap_index API)
self._target_body = self._model.get_body("target")
self._target_link_id = self._model.get_link_index("target")
assert self._target_body.is_mocap, "Target must be a mocap body"
# Initial cube position (in hand)
self._in_hand_pos = np.array(cfg.cube_initial_pos, dtype=np.float32)
@property
def observation_space(self):
return self._observation_space
@property
def action_space(self):
return self._action_space
def _extract_cube_states(self, data: mtx.SceneData, body: mtx.Body):
return body.get_position(data), body.get_rotation(data), data.dof_vel[:, body.get_dof_vel_indices()]
def _extract_link_states(self, data: mtx.SceneData, link_ids):
"""
Extract position, quaternion, and velocity for specified links.
Args:
data: SceneData object
link_ids: List of link indices or single int
Returns:
Tuple of (positions, quaternions, velocities)
- positions: (batch, num_links, 3)
- quaternions: (batch, num_links, 4) in (x, y, z, w) format
- velocities: (batch, num_links, 6) [linear_vel, angular_vel]
"""
# Ensure link_ids is a list
if isinstance(link_ids, int):
link_ids = [link_ids]
# Get all link poses: shape (batch, num_links_total, 7) [x, y, z, qx, qy, qz, qw]
all_poses = self._model.get_link_poses(data)
# Extract poses for requested links
poses = all_poses[:, link_ids, :] # (batch, num_requested_links, 7)
# Split into position and quaternion
positions = poses[:, :, :3]
quaternions = poses[:, :, 3:] # (qx, qy, qz, qw)
# TODO: Velocity computation - MotrixSim doesn't expose link velocities yet
# Temporary solution: use zero velocities
# Future options:
# 1. Finite differences (requires storing previous poses)
# 2. Compute from DOF velocities via Jacobian (if available)
# 3. Wait for API: model.get_link_velocities(data)
batch_size = data.shape[0]
num_links = len(link_ids)
velocities = np.zeros((batch_size, num_links, 6), dtype=np.float32)
for j in np.arange(num_links):
link = self._model.get_link(link_ids[j])
velocities[:, j] = np.concatenate(
(link.get_linear_velocity(data), link.get_angular_velocity(data)), axis=-1
)
return positions, quaternions, velocities
def apply_action(self, actions: np.ndarray, state: NpEnvState):
"""Apply actions to the hand actuators."""
cfg = self._cfg
# Scale actions from [-1, 1] to actuator control range
targets = utils.scale(actions, self._actuator_ctrl_lower, self._actuator_ctrl_upper)
# Apply action moving average for smoothness
if cfg.act_moving_average < 1.0:
targets = cfg.act_moving_average * targets + (1.0 - cfg.act_moving_average) * state.info["prev_actions"]
# Clamp to control limits
targets = np.clip(targets, self._actuator_ctrl_lower, self._actuator_ctrl_upper)
# Set actuator controls
state.data.actuator_ctrls = targets
state.info["prev_actions"] = targets.copy()
return state
def update_state(self, state: NpEnvState):
"""Update observations, rewards, and termination conditions."""
data = state.data
info = state.info
# compute obs
obs = self._compute_observation(state.data, info)
# Compute reward and termination
reward, terminated, goal_reached = self._compute_reward(state, info)
if np.any(goal_reached):
reset_goal_indices = np.where(goal_reached)[0]
self._reset_goal_pose(info, reset_goal_indices)
# Update target visualization
self._update_target_visualization(data, info)
state.obs = obs
state.reward = reward
state.terminated = terminated
return state
def _compute_observation(self, data: mtx.SceneData, info: dict):
cfg = self._cfg
num_envs = data.shape[0]
# Get hand DOF states
hand_dof_pos = data.dof_pos[:, : self._num_hand_dofs]
hand_dof_vel = data.dof_vel[:, : self._num_hand_dofs]
# Get cube state using link poses
cube_pos, cube_quat, cube_vel = self._extract_cube_states(data, self._cube_body)
cube_linvel = cube_vel[:, :3]
cube_angvel = cube_vel[:, 3:]
# Get fingertip states using link poses
fingertip_pos, fingertip_quat, fingertip_vel = self._extract_link_states(data, self._fingertip_link_ids)
# Flatten fingertip states (5 × 13 = 65)
fingertip_state = np.concatenate(
[
fingertip_pos.reshape(num_envs, -1), # 15
fingertip_quat.reshape(num_envs, -1), # 20
fingertip_vel.reshape(num_envs, -1), # 30
],
axis=-1,
) # Total: 65
# Compute relative quaternion
relative_quat = quaternion.mul(cube_quat, quaternion.conjugate(info["goal_rot"]))
scaled_hand_pos = utils.unscale(hand_dof_pos, self._hand_dof_lower_limits, self._hand_dof_upper_limits)
# Build observation (157 dims)
return np.concatenate(
[
scaled_hand_pos,
cfg.vel_obs_scale * hand_dof_vel, # 24
cube_pos, # 3
cube_quat, # 4
cube_linvel, # 3
cfg.vel_obs_scale * cube_angvel, # 3
info["goal_pos"], # 3
info["goal_rot"], # 4
relative_quat, # 4
fingertip_state, # 65
info["prev_actions"], # 20
],
axis=-1,
)
def _compute_reward(self, state: NpEnvState, info: dict):
"""
Reward components (3 core items):
1. Position distance penalty
2. Rotation alignment reward
3. Action regularization penalty
Additional rewards/penalties:
- Success bonus when goal is reached
- Fall penalty when cube drops
- Timeout penalty when episode ends without success
"""
cfg = self._cfg
num_envs = self._num_envs
# Get cube state using link poses
cube_pos, cube_quat, _ = self._extract_cube_states(state.data, self._cube_body)
# Distance from cube to goal position
goal_dist = np.linalg.norm(cube_pos - state.info["goal_pos"], axis=-1)
# Rotation distance
rot_dist = quaternion.rotation_distance(cube_quat, state.info["goal_rot"])
# Core reward components
dist_rew = goal_dist * cfg.dist_reward_scale
rot_rew = 1.0 / (np.abs(rot_dist) + cfg.rot_eps) * cfg.rot_reward_scale
action_penalty = np.sum(state.info["prev_actions"] ** 2, axis=-1) * cfg.action_penalty_scale
# Base reward
reward = dist_rew + rot_rew + action_penalty
# Check for success (only rotation tolerance)
goal_reached = np.abs(rot_dist) <= cfg.success_tolerance
# Update success counter
info["successes"] += goal_reached * 1
# Success bonus
reward = np.where(goal_reached, reward + cfg.reach_goal_bonus, reward)
# Fall penalty
fallen = goal_dist >= cfg.fall_dist
reward = np.where(fallen, reward + cfg.fall_penalty, reward)
# Termination conditions
terminated = np.zeros(num_envs, dtype=bool)
# 1. Fall termination
terminated = np.logical_or(terminated, fallen)
# 2. Success termination with hold mechanism
if cfg.max_consecutive_successes > 0:
# Reset progress on goal reached when max consecutive successes reached
new_pos = info["successes"] >= cfg.max_consecutive_successes
info["successes"] *= 1 - new_pos
# 3. NaN protection
terminated = np.logical_or(terminated, np.isnan(rot_dist))
terminated = np.logical_or(terminated, np.isnan(goal_dist))
return reward, terminated, new_pos
def _update_target_visualization(self, data: mtx.SceneData, info: dict):
"""Update the target mocap body to visualize the goal pose."""
cfg = self._cfg
# Compute visualization position (offset from goal position)
viz_pos = info["goal_pos"] + np.array(cfg.viz_target_offset, dtype=np.float32)
# Combine into pose array: [x, y, z, qx, qy, qz, qw]
viz_pose = np.concatenate([viz_pos, info["goal_rot"]], axis=-1)
# Update mocap body pose using correct API
self._target_body.mocap.set_pose(data, viz_pose)
def reset(self, data: mtx.SceneData):
"""Reset environments."""
cfg = self._cfg
# data is already filtered to contain only envs that need reset
num_resets = data.shape[0]
# Reset scene data
data.reset(self._model)
# Reset hand DOFs with noise
init_dof_pos = self._model.compute_init_dof_pos()
init_dof_vel = np.zeros(self._model.num_dof_vel, dtype=np.float32)
# Add noise to DOF positions
dof_pos_noise = np.random.uniform(
-cfg.reset_dof_pos_noise,
cfg.reset_dof_pos_noise,
(num_resets, self._num_hand_dofs),
)
# Add noise to DOF velocities
dof_vel_noise = np.random.uniform(
-cfg.reset_dof_vel_noise, cfg.reset_dof_vel_noise, (num_resets, self._num_hand_dofs)
)
# Set DOF states for all envs in data (already filtered)
dof_pos = np.tile(init_dof_pos, (num_resets, 1))
dof_vel = np.tile(init_dof_vel, (num_resets, 1))
dof_pos[:, : self._num_hand_dofs] += dof_pos_noise
dof_vel[:, : self._num_hand_dofs] += dof_vel_noise
data.set_dof_pos(dof_pos, self._model)
data.set_dof_vel(dof_vel)
# Reset cube position with small noise
cube_pos_noise = np.random.uniform(-cfg.reset_position_noise, cfg.reset_position_noise, (num_resets, 3))
cube_pos = np.tile(self._in_hand_pos, (num_resets, 1))
cube_pos += cube_pos_noise
# Randomize cube orientation
cube_quat = quaternion.generate_random_shoemake(num_resets)
# Set cube pose using body's set_dof_pos method
# Combine into DOF pose: [x, y, z, qx, qy, qz, qw]
cube_dof_pos = np.concatenate([cube_pos, cube_quat], axis=-1)
# Set cube DOF position
self._cube_body.set_dof_pos(data, cube_dof_pos)
# Set cube DOF velocity to zero
cube_dof_vel = np.zeros((num_resets, 6), dtype=np.float32) # 6DOF velocity
self._cube_body.set_dof_vel(data, cube_dof_vel)
# Reset goal pose
# Note: goal_pos and goal_rot are indexed by original env indices
info = {
"goal_pos": np.tile(self._in_hand_pos, num_resets).reshape(num_resets, 3),
"goal_rot": quaternion.generate_random_shoemake(num_resets),
"prev_actions": np.zeros((num_resets, self._num_actuators), dtype=np.float32),
"successes": np.zeros((num_resets), dtype=np.int32),
}
obs = self._compute_observation(data, info)
return obs, info
def _reset_goal_pose(self, info, env_ids):
"""Reset goal pose to random orientation with fixed position."""
num_resets = len(env_ids)
# Goal position is fixed
# Randomize goal orientation using Shoemake method for uniform SO(3) sampling
info["goal_rot"][env_ids] = quaternion.generate_random_shoemake(num_resets)

View File

@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2022 Shadow Robot Company Ltd
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.

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -0,0 +1,66 @@
<?xml version="1.0" encoding="utf-8"?>
<mujoco model="shadow_hand_cube_repose">
<include file="right_hand.xml" />
<compiler angle="radian" meshdir="assets" />
<option timestep="0.001" iterations="30" tolerance="1e-6"
cone="pyramidal" impratio="10" integrator="implicitfast">
<flag warmstart="enable" />
</option>
<size nconmax="800" njmax="2000" />
<statistic extent="0.4" center="0.3 0 0.15" />
<visual>
<rgba haze="0.15 0.25 0.35 1" />
<quality shadowsize="4096" />
<global azimuth="220" elevation="-30" />
</visual>
<default>
<default class="cube_contact">
<geom solimp="0.95 0.99 1e-3 0.5 2" solref="0.02 1"
friction="1.2 0.005 0.0001" condim="3" margin="0.003" />
</default>
</default>
<asset>
<texture type="2d" colorspace="sRGB" name="Cube_Side_Tex" file="assets/Cube_Side.png" />
<material name="cube" rgba="1 1 1 1" reflectance="0.5" metallic="0.0" roughness="1.0">
<layer texture="Cube_Side_Tex" role="rgb" />
</material>
<material name="cube_target" rgba="1 1 1 0.7" reflectance="0.5" metallic="0.0" roughness="1.0">
<layer texture="Cube_Side_Tex" role="rgb" />
</material>
<mesh name="cube_vis_mesh" file="Cube_Side.obj" scale="0.8333 0.8333 0.8333" />
<texture name="motphys-ground" type="2d" file="../../../common/motphys-ground.png" />
<material name="motphys-ground" texture="motphys-ground" texuniform="true" texrepeat="0.4 0.4" />
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom type="plane" size="0 0 .01" material="motphys-ground"/>
<body name="cube" pos="0.30 0.00 0.35">
<freejoint name="cube_joint" />
<geom name="cube_vis" mesh="cube_vis_mesh" type="mesh" size="0.025 0.025 0.025"
class="cube_contact" material="cube" contype="0" conaffinity="0"
group="1" />
<geom type="box" size="0.025 0.025 0.025"
group="3"
density="225"
rgba="0.8 0.2 0.2 0.5" />
</body>
<body name="target" pos="0.45 -0.20 0.22" mocap="true">
<geom name="target_cube" mesh="cube_vis_mesh" type="mesh" size="0.025 0.025 0.025"
contype="0" conaffinity="0"
material="cube_target" />
</body>
</worldbody>
</mujoco>

View File

@@ -0,0 +1,317 @@
<mujoco model="right_shadow_hand">
<compiler angle="radian" meshdir="assets" autolimits="true"/>
<option cone="elliptic" impratio="10"/>
<default>
<default class="right_hand">
<mesh scale="0.001 0.001 0.001"/>
<joint axis="1 0 0" damping="0.05" armature="0.005" frictionloss="0.01"/>
<position forcerange="-1 1"/>
<default class="wrist">
<joint damping="0.5"/>
<default class="wrist_y">
<joint axis="0 1 0" range="-0.523599 0.174533"/>
<position kp="10" ctrlrange="-0.523599 0.174533" forcerange="-10 10"/>
</default>
<default class="wrist_x">
<joint range="-0.698132 0.488692"/>
<position kp="8" ctrlrange="-0.698132 0.488692" forcerange="-5 5"/>
</default>
</default>
<default class="thumb">
<default class="thbase">
<joint axis="0 0 -1" range="-1.0472 1.0472"/>
<position kp="0.4" ctrlrange="-1.0472 1.0472" forcerange="-3 3"/>
</default>
<default class="thproximal">
<joint range="0 1.22173"/>
<position ctrlrange="0 1.22173" forcerange="-2 2"/>
</default>
<default class="thhub">
<joint range="-0.20944 0.20944"/>
<position kp="0.5" ctrlrange="-0.20944 0.20944"/>
</default>
<default class="thmiddle">
<joint axis="0 -1 0" range="-0.698132 0.698132"/>
<position kp="1.5" ctrlrange="-0.698132 0.698132"/>
</default>
<default class="thdistal">
<joint range="-0.261799 1.5708"/>
<position ctrlrange="-0.261799 1.5708"/>
</default>
</default>
<default class="metacarpal">
<joint axis="0.573576 0 0.819152" range="0 0.785398"/>
<position ctrlrange="0 0.785398"/>
</default>
<default class="knuckle">
<joint axis="0 -1 0" range="-0.349066 0.349066"/>
<position ctrlrange="-0.349066 0.349066"/>
</default>
<default class="proximal">
<joint range="-0.261799 1.5708"/>
<position ctrlrange="-0.261799 1.5708"/>
</default>
<default class="middle_distal">
<joint range="0 1.5708"/>
<position kp="0.5" ctrlrange="0 3.1415"/>
</default>
<default class="plastic">
<geom solimp="0.5 0.99 0.0001" solref="0.005 1"/>
<default class="plastic_visual">
<geom type="mesh" material="black" contype="0" conaffinity="0" group="2"/>
</default>
<default class="plastic_collision">
<geom group="3"/>
</default>
</default>
</default>
</default>
<asset>
<material name="black" specular="0.5" shininess="0.25" rgba="0.16355 0.16355 0.16355 1"/>
<material name="gray" specular="0.0" shininess="0.25" rgba="0.80848 0.80848 0.80848 1"/>
<material name="metallic" specular="0" shininess="0.25" rgba="0.9 0.9 0.9 1"/>
<mesh class="right_hand" file="forearm_0.obj"/>
<mesh class="right_hand" file="forearm_1.obj"/>
<mesh class="right_hand" file="forearm_collision.obj"/>
<mesh class="right_hand" file="wrist.obj"/>
<mesh class="right_hand" file="palm.obj"/>
<mesh class="right_hand" file="f_knuckle.obj"/>
<mesh class="right_hand" file="f_proximal.obj"/>
<mesh class="right_hand" file="f_middle.obj"/>
<mesh class="right_hand" file="f_distal_pst.obj"/>
<mesh class="right_hand" file="lf_metacarpal.obj"/>
<mesh class="right_hand" file="th_proximal.obj"/>
<mesh class="right_hand" file="th_middle.obj"/>
<mesh class="right_hand" file="th_distal_pst.obj"/>
</asset>
<worldbody>
<body name="rh_forearm" childclass="right_hand" quat="0 1 0 1" pos="0 0 0.2">
<inertial mass="3" pos="0 0 0.09" diaginertia="0.0138 0.0138 0.00744"/>
<geom class="plastic_visual" mesh="forearm_0" material="gray"/>
<geom class="plastic_visual" mesh="forearm_1"/>
<geom class="plastic_collision" type="mesh" mesh="forearm_collision"/>
<geom class="plastic_collision" size="0.035 0.035 0.035" pos="0.01 0.0 0.181" quat="0.380188 0.924909 0 0"
type="box"/>
<body name="rh_wrist" pos="0.01 0 0.21301" quat="1 0 0 1">
<inertial mass="0.1" pos="0 0 0.029" quat="0.5 0.5 0.5 0.5" diaginertia="6.4e-05 4.38e-05 3.5e-05"/>
<joint class="wrist_y" name="rh_WRJ2"/>
<geom class="plastic_visual" mesh="wrist" material="metallic"/>
<geom size="0.0135 0.015" quat="0.499998 0.5 0.5 -0.500002" type="cylinder" class="plastic_collision"/>
<geom size="0.011 0.005" pos="-0.026 0 0.034" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<geom size="0.011 0.005" pos="0.031 0 0.034" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<geom size="0.0135 0.009 0.005" pos="-0.021 0 0.011" quat="0.923879 0 0.382684 0" type="box"
class="plastic_collision"/>
<geom size="0.0135 0.009 0.005" pos="0.026 0 0.01" quat="0.923879 0 -0.382684 0" type="box"
class="plastic_collision"/>
<body name="rh_palm" pos="0 0 0.034">
<inertial mass="0.3" pos="0 0 0.035" quat="1 0 0 1" diaginertia="0.0005287 0.0003581 0.000191"/>
<joint class="wrist_x" name="rh_WRJ1"/>
<site name="grasp_site" pos="0 -.035 0.09" group="4"/>
<geom class="plastic_visual" mesh="palm"/>
<geom size="0.031 0.0035 0.049" pos="0.011 0.0085 0.038" type="box" class="plastic_collision"/>
<geom size="0.018 0.0085 0.049" pos="-0.002 -0.0035 0.038" type="box" class="plastic_collision"/>
<geom size="0.013 0.0085 0.005" pos="0.029 -0.0035 0.082" type="box" class="plastic_collision"/>
<geom size="0.013 0.007 0.009" pos="0.0265 -0.001 0.07" quat="0.987241 0.0990545 0.0124467 0.124052"
type="box" class="plastic_collision"/>
<geom size="0.0105 0.0135 0.012" pos="0.0315 -0.0085 0.001" type="box" class="plastic_collision"/>
<geom size="0.011 0.0025 0.015" pos="0.0125 -0.015 0.004" quat="0.971338 0 0 -0.237703" type="box"
class="plastic_collision"/>
<geom size="0.009 0.012 0.002" pos="0.011 0 0.089" type="box" class="plastic_collision"/>
<geom size="0.01 0.012 0.02" pos="-0.03 0 0.009" type="box" class="plastic_collision"/>
<body name="rh_ffknuckle" pos="0.033 0 0.095">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_FFJ4" class="knuckle"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<body name="rh_ffproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_FFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision"/>
<body name="rh_ffmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_FFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision"/>
<body name="rh_ffdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_FFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="capsule" pos="0 0 0.012" size="0.00705 0.012"/>
</body>
</body>
</body>
</body>
<body name="rh_mfknuckle" pos="0.011 0 0.099">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_MFJ4" class="knuckle"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<body name="rh_mfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_MFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision"/>
<body name="rh_mfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_MFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision"/>
<body name="rh_mfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_MFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="capsule" pos="0 0 0.012" size="0.00705 0.012"/>
</body>
</body>
</body>
</body>
<body name="rh_rfknuckle" pos="-0.011 0 0.095">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_RFJ4" class="knuckle" axis="0 1 0"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<body name="rh_rfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_RFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision"/>
<body name="rh_rfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_RFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision"/>
<body name="rh_rfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_RFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="capsule" pos="0 0 0.012" size="0.00705 0.012"/>
</body>
</body>
</body>
</body>
<body name="rh_lfmetacarpal" pos="-0.033 0 0.02071">
<inertial mass="0.03" pos="0 0 0.04" quat="1 0 0 1" diaginertia="1.638e-05 1.45e-05 4.272e-06"/>
<joint name="rh_LFJ5" class="metacarpal"/>
<geom class="plastic_visual" mesh="lf_metacarpal"/>
<geom size="0.011 0.012 0.025" pos="0.002 0 0.033" type="box" class="plastic_collision"/>
<body name="rh_lfknuckle" pos="0 0 0.06579">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_LFJ4" class="knuckle" axis="0 1 0"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision"/>
<body name="rh_lfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_LFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision"/>
<body name="rh_lfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_LFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision"/>
<body name="rh_lfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_LFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="capsule" pos="0 0 0.012" size="0.00705 0.012"/>
</body>
</body>
</body>
</body>
</body>
<body name="rh_thbase" pos="0.034 -0.00858 0.029" quat="0.92388 0 0.382683 0">
<inertial mass="0.01" pos="0 0 0" diaginertia="1.6e-07 1.6e-07 1.6e-07"/>
<joint name="rh_THJ5" class="thbase"/>
<geom class="plastic_collision" size="0.013"/>
<body name="rh_thproximal">
<inertial mass="0.04" pos="0 0 0.019" diaginertia="1.36e-05 1.36e-05 3.13e-06"/>
<joint name="rh_THJ4" class="thproximal"/>
<geom class="plastic_visual" mesh="th_proximal"/>
<geom class="plastic_collision" size="0.0105 0.009" pos="0 0 0.02" type="capsule"/>
<body name="rh_thhub" pos="0 0 0.038">
<inertial mass="0.005" pos="0 0 0" diaginertia="1e-06 1e-06 3e-07"/>
<joint name="rh_THJ3" class="thhub"/>
<geom size="0.011" class="plastic_collision"/>
<body name="rh_thmiddle">
<inertial mass="0.02" pos="0 0 0.016" diaginertia="5.1e-06 5.1e-06 1.21e-06"/>
<joint name="rh_THJ2" class="thmiddle"/>
<geom class="plastic_visual" mesh="th_middle"/>
<geom size="0.009 0.009" pos="0 0 0.012" type="capsule" class="plastic_collision"/>
<geom size="0.01" pos="0 0 0.03" class="plastic_collision"/>
<body name="rh_thdistal" pos="0 0 0.032" quat="1 0 0 -1">
<inertial mass="0.017" pos="0 0 0.0145588" quat="1 0 0 1"
diaginertia="2.37794e-06 2.27794e-06 1e-06"/>
<joint name="rh_THJ1" class="thdistal"/>
<geom class="plastic_visual" mesh="th_distal_pst"/>
<geom class="plastic_collision" type="capsule" pos="0 0 0.013" size="0.00918 0.013"/>
</body>
</body>
</body>
</body>
</body>
</body>
</body>
</body>
</worldbody>
<contact>
<exclude body1="rh_wrist" body2="rh_forearm"/>
<exclude body1="rh_thproximal" body2="rh_thmiddle"/>
</contact>
<tendon>
<fixed name="rh_FFJ0">
<joint joint="rh_FFJ2" coef="1"/>
<joint joint="rh_FFJ1" coef="1"/>
</fixed>
<fixed name="rh_MFJ0">
<joint joint="rh_MFJ2" coef="1"/>
<joint joint="rh_MFJ1" coef="1"/>
</fixed>
<fixed name="rh_RFJ0">
<joint joint="rh_RFJ2" coef="1"/>
<joint joint="rh_RFJ1" coef="1"/>
</fixed>
<fixed name="rh_LFJ0">
<joint joint="rh_LFJ2" coef="1"/>
<joint joint="rh_LFJ1" coef="1"/>
</fixed>
</tendon>
<actuator>
<position name="rh_A_WRJ2" joint="rh_WRJ2" class="wrist_y"/>
<position name="rh_A_WRJ1" joint="rh_WRJ1" class="wrist_x"/>
<position name="rh_A_THJ5" joint="rh_THJ5" class="thbase"/>
<position name="rh_A_THJ4" joint="rh_THJ4" class="thproximal"/>
<position name="rh_A_THJ3" joint="rh_THJ3" class="thhub"/>
<position name="rh_A_THJ2" joint="rh_THJ2" class="thmiddle"/>
<position name="rh_A_THJ1" joint="rh_THJ1" class="thdistal"/>
<position name="rh_A_FFJ4" joint="rh_FFJ4" class="knuckle"/>
<position name="rh_A_FFJ3" joint="rh_FFJ3" class="proximal"/>
<position name="rh_A_FFJ0" tendon="rh_FFJ0" class="middle_distal"/>
<position name="rh_A_MFJ4" joint="rh_MFJ4" class="knuckle"/>
<position name="rh_A_MFJ3" joint="rh_MFJ3" class="proximal"/>
<position name="rh_A_MFJ0" tendon="rh_MFJ0" class="middle_distal"/>
<position name="rh_A_RFJ4" joint="rh_RFJ4" class="knuckle"/>
<position name="rh_A_RFJ3" joint="rh_RFJ3" class="proximal"/>
<position name="rh_A_RFJ0" tendon="rh_RFJ0" class="middle_distal"/>
<position name="rh_A_LFJ5" joint="rh_LFJ5" class="metacarpal"/>
<position name="rh_A_LFJ4" joint="rh_LFJ4" class="knuckle"/>
<position name="rh_A_LFJ3" joint="rh_LFJ3" class="proximal"/>
<position name="rh_A_LFJ0" tendon="rh_LFJ0" class="middle_distal"/>
</actuator>
</mujoco>

View File

@@ -12,139 +12,290 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
from . import utils
class Quaternion:
def multiply(q1, q2):
"""
Quaternion multiply, with [w, x, y, z] format
"""
qx1, qy1, qz1, qw1 = q1[0], q1[1], q1[2], q1[3]
qx2, qy2, qz2, qw2 = q2[0], q2[1], q2[2], q2[3]
qw = qw1 * qw2 - qx1 * qx2 - qy1 * qy2 - qz1 * qz2
qx = qw1 * qx2 + qx1 * qw2 + qy1 * qz2 - qz1 * qy2
qy = qw1 * qy2 - qx1 * qz2 + qy1 * qw2 + qz1 * qx2
qz = qw1 * qz2 + qx1 * qy2 - qy1 * qx2 + qz1 * qw2
def mul(q1, q2):
"""
Multiply two quaternions.
return np.array([qx, qy, qz, qw], dtype=np.float32)
Quaternion format: (x, y, z, w)
def from_euler(roll: np.ndarray, pitch: np.ndarray, yaw: np.ndarray):
"""
Euler convert to quaternion, with [w, x, y, z] format
"""
cy = np.cos(yaw * 0.5)
sy = np.sin(yaw * 0.5)
cp = np.cos(pitch * 0.5)
sp = np.sin(pitch * 0.5)
cr = np.cos(roll * 0.5)
sr = np.sin(roll * 0.5)
Args:
q1: First quaternion(s). Shape: (..., 4)
q2: Second quaternion(s). Shape: (..., 4)
qw = cr * cp * cy + sr * sp * sy
qx = sr * cp * cy - cr * sp * sy
qy = cr * sp * cy + sr * cp * sy
qz = cr * cp * sy - sr * sp * cy
Returns:
Product quaternion(s). Shape: (..., 4)
"""
x1, y1, z1, w1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3]
x2, y2, z2, w2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3]
return np.stack([qx, qy, qz, qw], dtype=np.float32, axis=-1)
# Standard quaternion multiplication formula
w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
def rotate_vector(quats: np.ndarray, v: np.ndarray):
"""
Rotate a list vectors v by a list of quaternions using a vectorized approach. v could be a simple shape (3,)
vector, or a shape (N,3) vector array with quats shape (N,4)
return np.stack([x, y, z, w], axis=-1)
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
def conjugate(q):
"""
Compute the conjugate of a quaternion.
# Extract the scalar (w) and vector (x, y, z) parts of the quaternions
w = quats[:, -1] # Shape (N,)
im = quats[:, :3] # Shape (N, 3)
For q = (x, y, z, w), conjugate is (-x, -y, -z, w)
t = 2 * np.cross(im, v)
return v + w.reshape(-1, 1) * t + np.cross(im, t)
Args:
q: Input quaternion(s). Shape: (..., 4)
def rotate_inverse(quats, v):
"""
Rotate a list of vectors v by a list of inverse quaternions using a vectorized approach.
Returns:
Conjugate quaternion(s). Shape: (..., 4)
"""
return q * np.array([-1, -1, -1, 1], dtype=q.dtype)
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
def from_euler(roll: np.ndarray, pitch: np.ndarray, yaw: np.ndarray):
"""
Euler convert to quaternion, with [x, y, z, w] format
"""
cy = np.cos(yaw * 0.5)
sy = np.sin(yaw * 0.5)
cp = np.cos(pitch * 0.5)
sp = np.sin(pitch * 0.5)
cr = np.cos(roll * 0.5)
sr = np.sin(roll * 0.5)
# Extract the scalar (w) and vector (x, y, z) parts of the quaternions
w = quats[:, -1] # Shape (N,)
im = quats[:, :3] # Shape (N, 3)
qw = cr * cp * cy + sr * sp * sy
qx = sr * cp * cy - cr * sp * sy
qy = cr * sp * cy + sr * cp * sy
qz = cr * cp * sy - sr * sp * cy
# 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)
return np.stack([qx, qy, qz, qw], dtype=np.float32, axis=-1)
term = cross_im_v - w.reshape(-1, 1) * v
# Final result: v' = v + 2 * r × term
v_rotated = v + 2 * np.cross(im, term)
def from_angle_axis(angle, axis):
"""
Create quaternion from angle-axis representation.
return v_rotated
Args:
angle: Rotation angle in radians. Shape: (batch,) or scalar
axis: Rotation axis (will be normalized). Shape: (batch, 3) or (3,)
def similarity(q_current, q_target):
"""
Use NumPy to compute attitude alignment reward between two batches of quaternions.
Returns:
Quaternion in (x, y, z, w) format. Shape: (batch, 4) or (4,)
"""
# Ensure angle has proper shape for broadcasting
if np.isscalar(angle):
angle = np.array([angle])
Parameters:
q_current (np.ndarray): Quaternion of current pose, shape (num_envs, 4).
q_target (np.ndarray): Quaternion of target pose, shape (num_envs, 4) or (4,).
If (4,), it will be broadcast to all environments.
# Normalize axis
axis = utils.normalize(axis)
Returns:
np.ndarray: Reward value for each environment, shape (num_envs,). Reward value range is [-1, 1].
"""
# Ensure input is float array
q_current = q_current.astype(np.float32)
q_target = q_target.astype(np.float32)
# Compute half angle
half_angle = angle / 2.0
# If q_target is a single quaternion, broadcast to all environments
if q_target.ndim == 1:
# Use np.tile for broadcasting
q_target = np.tile(q_target, (q_current.shape[0], 1))
# Compute quaternion components
sin_half = np.sin(half_angle)
cos_half = np.cos(half_angle)
# Step 1: Compute conjugate of q_current
# Conjugate of quaternion (w, x, y, z) is (w, -x, -y, -z)
q_current_conj = np.copy(q_current)
q_current_conj[..., 1:] *= -1 # Negate x, y, z components
# Handle broadcasting
if axis.ndim == 1:
# Single axis
w = cos_half
xyz = axis * sin_half[..., np.newaxis]
else:
# Multiple axes
w = cos_half
xyz = axis * sin_half[..., np.newaxis]
# Step 2: Compute relative quaternion q_rel = q_target * q_current_conj
# Unpack quaternion components for computation
w1, x1, y1, z1 = q_target[..., 0], q_target[..., 1], q_target[..., 2], q_target[..., 3]
w2, x2, y2, z2 = q_current_conj[..., 0], q_current_conj[..., 1], q_current_conj[..., 2], q_current_conj[..., 3]
return utils.normalize(np.concatenate([xyz, w[..., np.newaxis]], axis=-1))
# Apply quaternion multiplication formula
w_rel = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
# For numerical stability, clamp w_rel to [-1.0, 1.0] range
w_rel_clamped = np.clip(w_rel, -1.0, 1.0)
def rotate_vector(quats: np.ndarray, v: np.ndarray):
"""
Rotate a list vectors v by a list of quaternions using a vectorized approach. v could be a simple shape (3,)
vector, or a shape (N,3) vector array with quats shape (N,4)
# Step 3: Compute rotation angle theta
theta = 2.0 * np.arccos(w_rel_clamped)
Parameters:
quats (np.ndarray): Array of quaternions of shape (N, 4). Each quaternion is in [x, y, z, w] format.
v (np.ndarray): Fixed vector of shape (3,) to be rotated.
# Step 4: Compute reward
reward = np.cos(theta)
Returns:
np.ndarray: Array of rotated vectors of shape (N, 3).
"""
# Normalize the quaternions to ensure they are unit quaternions
return reward
# Extract the scalar (w) and vector (x, y, z) parts of the quaternions
w = quats[:, -1] # Shape (N,)
im = quats[:, :3] # Shape (N, 3)
def get_yaw(quat: np.ndarray) -> np.ndarray:
qx, qy, qz, qw = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
# Compute yaw angle (rotation around Z axis)
siny_cosp = 2 * (qw * qz + qx * qy)
cosy_cosp = 1 - 2 * (qy * qy + qz * qz)
return np.arctan2(siny_cosp, cosy_cosp)
t = 2 * np.cross(im, v)
return v + w.reshape(-1, 1) * t + np.cross(im, t)
def rotate_inverse(quats, v):
"""
Rotate a list of vectors v by a list of inverse quaternions using a vectorized approach.
Parameters:
quats (np.ndarray): Array of quaternions of shape (N, 4). Each quaternion is in [x, y, z, w] 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)
term = cross_im_v - w.reshape(-1, 1) * v
# Final result: v' = v + 2 * r × term
v_rotated = v + 2 * np.cross(im, term)
return v_rotated
def similarity(q_current, q_target):
"""
Use NumPy to compute attitude alignment reward between two batches of quaternions.
Parameters:
q_current (np.ndarray): Quaternion of current pose, shape (num_envs, 4).
q_target (np.ndarray): Quaternion of target pose, shape (num_envs, 4) or (4,).
If (4,), it will be broadcast to all environments.
Returns:
np.ndarray: Reward value for each environment, shape (num_envs,). Reward value range is [-1, 1].
"""
# Ensure input is float array
q_current = q_current.astype(np.float32)
q_target = q_target.astype(np.float32)
# If q_target is a single quaternion, broadcast to all environments
if q_target.ndim == 1:
# Use np.tile for broadcasting
q_target = np.tile(q_target, (q_current.shape[0], 1))
# Step 1: Compute conjugate of q_current
# Conjugate of quaternion (x, y, z, w) is (-x, -y, -z, w)
q_current_conj = np.copy(q_current)
q_current_conj[..., :3] *= -1 # Negate x, y, z components
# Step 2: Compute relative quaternion q_rel = q_target * q_current_conj
# Unpack quaternion components for computation
x1, y1, z1, w1 = q_target[..., 0], q_target[..., 1], q_target[..., 2], q_target[..., 3]
x2, y2, z2, w2 = q_current_conj[..., 0], q_current_conj[..., 1], q_current_conj[..., 2], q_current_conj[..., 3]
# Apply quaternion multiplication formula
w_rel = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
# For numerical stability, clamp w_rel to [-1.0, 1.0] range
w_rel_clamped = np.clip(w_rel, -1.0, 1.0)
# Step 3: Compute rotation angle theta
theta = 2.0 * np.arccos(w_rel_clamped)
# Step 4: Compute reward
reward = np.cos(theta)
return reward
def rotation_distance(q1, q2):
"""
Compute the rotation distance between two quaternions in radians.
Args:
q1: First quaternion(s). Shape: (..., 4)
q2: Second quaternion(s). Shape: (..., 4)
Returns:
Rotation distance in radians. Shape: (...)
"""
quat_diff = mul(q1, conjugate(q2))
# Extract imaginary part (x, y, z) and compute norm
imaginary_norm = np.linalg.norm(quat_diff[..., :3], axis=-1)
# Clamp to valid range for arcsin
imaginary_norm = np.clip(imaginary_norm, 0.0, 1.0)
return 2.0 * np.arcsin(imaginary_norm)
def get_euler_xyz(q: np.ndarray):
"""
Convert quaternion to Euler angles (roll, pitch, yaw).
Args:
q: Quaternion(s) in (x, y, z, w) format. Shape: (..., 4)
Returns:
Tuple of (roll, pitch, yaw) in radians. Each has shape: (...)
"""
x, y, z, w = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
# Roll (x-axis rotation)
sinr_cosp = 2.0 * (w * x + y * z)
cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
roll = np.arctan2(sinr_cosp, cosr_cosp)
# Pitch (y-axis rotation)
sinp = 2.0 * (w * y - z * x)
pitch = np.where(np.abs(sinp) >= 1, np.copysign(np.pi / 2.0, sinp), np.arcsin(sinp))
# Yaw (z-axis rotation)
siny_cosp = 2.0 * (w * z + x * y)
cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
yaw = np.arctan2(siny_cosp, cosy_cosp)
return roll, pitch, yaw
def get_yaw(quat: np.ndarray) -> np.ndarray:
_, _, yaw = get_euler_xyz(quat)
return yaw
def generate_random_shoemake(size):
"""
Generate uniformly distributed random quaternions using Shoemake's method.
This ensures uniform distribution over SO(3) rotation space.
Args:
size: Number of quaternions to generate (int or tuple)
Returns:
Random quaternions in (x, y, z, w) format. Shape: (size, 4) or (*size, 4)
Reference:
K. Shoemake, "Uniform Random Rotations", Graphics Gems III, 1992
"""
# Convert size to tuple if it's a scalar (handles numpy int64)
if np.isscalar(size):
size = (int(size),)
elif not isinstance(size, tuple):
size = tuple(size)
# Generate three uniform random numbers
u1, u2, u3 = np.random.uniform(0, 1, size=(3, *size))
# Shoemake's method
sqrt1_u1 = np.sqrt(1 - u1)
sqrtu1 = np.sqrt(u1)
w = sqrt1_u1 * np.sin(2 * np.pi * u2)
x = sqrt1_u1 * np.cos(2 * np.pi * u2)
y = sqrtu1 * np.sin(2 * np.pi * u3)
z = sqrtu1 * np.cos(2 * np.pi * u3)
return np.stack([x, y, z, w], axis=-1)

View File

@@ -0,0 +1,65 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
def scale(x, lower, upper):
"""
Scale from normalized [-1, 1] range to [lower, upper] range.
Args:
x: Input in range [-1, 1]. Shape: (batch, dims) or (dims,)
lower: Lower bound. Shape: (dims,) or scalar
upper: Upper bound. Shape: (dims,) or scalar
Returns:
Scaled values in range [lower, upper]
"""
return 0.5 * (x + 1.0) * (upper - lower) + lower
def unscale(x, lower, upper):
"""
Scale from [lower, upper] range to normalized [-1, 1] range.
"""
rng = upper - lower
safe_rng = np.where(rng == 0, 1.0, rng)
res = 2.0 * (x - lower) / safe_rng - 1.0
if np.any(np.isnan(res)):
print("[DEBUG] math_utils.unscale: 产生 NaN!")
print(f" -> x: {x}")
print(f" -> lower: {lower}")
print(f" -> upper: {upper}")
print(f" -> range: {upper - lower}")
return np.where(rng == 0, 0.0, res)
def normalize(x):
"""
Normalize a vector to unit length.
Args:
x: Input vectors. Shape: (..., n)
eps: Minimum norm to avoid division by zero
Returns:
Normalized vectors. Shape: (..., n)
"""
norm = np.linalg.norm(x, axis=-1, keepdims=True)
if norm > 0.0:
return x / norm
else:
raise ValueError("Zero vector could not be normalized.")

View File

@@ -169,14 +169,12 @@ class NpEnv(ABEnv):
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