chore: release v0.2.0
This commit is contained in:
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
16
motrix_envs/src/motrix_envs/basic/acrobot/__init__.py
Normal file
16
motrix_envs/src/motrix_envs/basic/acrobot/__init__.py
Normal 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
|
||||
49
motrix_envs/src/motrix_envs/basic/acrobot/acrobot.xml
Normal file
49
motrix_envs/src/motrix_envs/basic/acrobot/acrobot.xml
Normal 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>
|
||||
152
motrix_envs/src/motrix_envs/basic/acrobot/acrobot_np.py
Normal file
152
motrix_envs/src/motrix_envs/basic/acrobot/acrobot_np.py
Normal 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
|
||||
33
motrix_envs/src/motrix_envs/basic/acrobot/cfg.py
Normal file
33
motrix_envs/src/motrix_envs/basic/acrobot/cfg.py
Normal 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
|
||||
@@ -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>
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
16
motrix_envs/src/motrix_envs/basic/finger/__init__.py
Normal file
16
motrix_envs/src/motrix_envs/basic/finger/__init__.py
Normal 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
|
||||
102
motrix_envs/src/motrix_envs/basic/finger/cfg.py
Normal file
102
motrix_envs/src/motrix_envs/basic/finger/cfg.py
Normal 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"
|
||||
76
motrix_envs/src/motrix_envs/basic/finger/finger.xml
Normal file
76
motrix_envs/src/motrix_envs/basic/finger/finger.xml
Normal 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>
|
||||
431
motrix_envs/src/motrix_envs/basic/finger/finger_np.py
Normal file
431
motrix_envs/src/motrix_envs/basic/finger/finger_np.py
Normal 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
|
||||
77
motrix_envs/src/motrix_envs/basic/finger/finger_spin.xml
Normal file
77
motrix_envs/src/motrix_envs/basic/finger/finger_spin.xml
Normal 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>
|
||||
|
||||
@@ -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>
|
||||
@@ -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>
|
||||
16
motrix_envs/src/motrix_envs/basic/humanoid/__init__.py
Normal file
16
motrix_envs/src/motrix_envs/basic/humanoid/__init__.py
Normal 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
|
||||
84
motrix_envs/src/motrix_envs/basic/humanoid/cfg.py
Normal file
84
motrix_envs/src/motrix_envs/basic/humanoid/cfg.py
Normal 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
|
||||
207
motrix_envs/src/motrix_envs/basic/humanoid/humanoid.xml
Normal file
207
motrix_envs/src/motrix_envs/basic/humanoid/humanoid.xml
Normal 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>
|
||||
550
motrix_envs/src/motrix_envs/basic/humanoid/humanoid_np.py
Normal file
550
motrix_envs/src/motrix_envs/basic/humanoid/humanoid_np.py
Normal 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,))
|
||||
16
motrix_envs/src/motrix_envs/basic/manipulator/__init__.py
Normal file
16
motrix_envs/src/motrix_envs/basic/manipulator/__init__.py
Normal 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
|
||||
95
motrix_envs/src/motrix_envs/basic/manipulator/cfg.py
Normal file
95
motrix_envs/src/motrix_envs/basic/manipulator/cfg.py
Normal 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
|
||||
@@ -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>
|
||||
579
motrix_envs/src/motrix_envs/basic/manipulator/manipulator_np.py
Normal file
579
motrix_envs/src/motrix_envs/basic/manipulator/manipulator_np.py
Normal 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)
|
||||
16
motrix_envs/src/motrix_envs/basic/pendulum/__init__.py
Normal file
16
motrix_envs/src/motrix_envs/basic/pendulum/__init__.py
Normal 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
|
||||
39
motrix_envs/src/motrix_envs/basic/pendulum/cfg.py
Normal file
39
motrix_envs/src/motrix_envs/basic/pendulum/cfg.py
Normal 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)))
|
||||
35
motrix_envs/src/motrix_envs/basic/pendulum/pendulum.xml
Normal file
35
motrix_envs/src/motrix_envs/basic/pendulum/pendulum.xml
Normal 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>
|
||||
|
||||
117
motrix_envs/src/motrix_envs/basic/pendulum/pendulum_np.py
Normal file
117
motrix_envs/src/motrix_envs/basic/pendulum/pendulum_np.py
Normal 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)}
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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>
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
142
motrix_envs/src/motrix_envs/manipulation/shadow_hand/cfg.py
Normal file
142
motrix_envs/src/motrix_envs/manipulation/shadow_hand/cfg.py
Normal 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
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d7d001736c76ac85f8ce89c8303ee489e1e8fa54969f52052b6a32108c8fefb6
|
||||
size 99120
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:da898c6b769ea0d673cbad5b40674643edd12444bce373e30eb109286646bfa9
|
||||
size 1127
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b2271d00d1aeb73da9554d04c1a5fd317039ea02d4bc1a6cb842bcdacc2c8599
|
||||
size 357939
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:350fcd33bc8ead915c0f646c8303d6b7554ad50ce546a128671d8d1fc75d5807
|
||||
size 341254
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:e997543c6c635fa80d152655869fe1d5bc1000a55f7fb648544a8755f068a075
|
||||
size 34351
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:8e78001a2ff735088cd71e009498ac73e4f966a6055d3bc864db049eacdbda11
|
||||
size 25738
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:ec7741fb5a9ee379faec68091de1086ade627b46a73f0a32b3cc09fec72524e3
|
||||
size 64846
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:4d175854bbd3c93627eb6b9f09296efd4e133cd72885b9bc57c82f902c8c94c4
|
||||
size 80439
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:6ada902a47f61f000fb14cbb8245e2093ff32e3e6b43deea790fe5d034864774
|
||||
size 1433615
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:38328466d52194888812a15da0dd3d5a1ef786a005121338f6a5a8ce3057602d
|
||||
size 26812
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:71d5317a6c0a5dcf6070791981f2a4720919d0419d0c17b400aef91343ce6baa
|
||||
size 88696
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:32897f8d46e463657a822943c31f440dc6bbac87b5c7939d10cdbf971291e8d9
|
||||
size 118690
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:614cf65374fafaceb713f4ee711692be6a3a0a7851ff2e720925f3a858cb3882
|
||||
size 775747
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:ca87b937114face76a91810f2ee60ea4fc217b57ffa9a7a7f166169a3ebf7da1
|
||||
size 300842
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a83bd4451f98d88c6c43ca0bbba30b6fd95b205da13adf9aff5a56c5aff3b4a8
|
||||
size 75080
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:71ffd5329c80a0cef890da0ff490b8e04200ff53154e5b416ce0be24a6a1b256
|
||||
size 23772
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:fc5620e4d9ea46d9526f0b40f8586bb5607e8ca763e5e23dab907a1baaf4890d
|
||||
size 110838
|
||||
@@ -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>
|
||||
@@ -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>
|
||||
@@ -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)
|
||||
|
||||
65
motrix_envs/src/motrix_envs/math/utils.py
Normal file
65
motrix_envs/src/motrix_envs/math/utils.py
Normal 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.")
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user