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

add go2 training

Browse Source
This commit is contained in:
jyf23
2026-01-02 21:10:34 +08:00
committed by motphys-developers
parent 62011bb24f
commit dbfa9e31fa
23 changed files with 952 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
motrix_envs/src/motrix_envs/locomotion/__init__.py
View File

@@ -13,4 +13,4 @@
# limitations under the License.
# ==============================================================================
from . import anymal_c, go1 # noqa: F401 register envs
from . import anymal_c, go1, go2 # noqa: F401 register envs

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

@@ -0,0 +1,16 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from . import walk_np # noqa: F401 register envs

139
motrix_envs/src/motrix_envs/locomotion/go2/cfg.py Normal file
View File

@@ -0,0 +1,139 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os
from dataclasses import dataclass, field
from motrix_envs import registry
from motrix_envs.base import EnvCfg
model_file = os.path.dirname(__file__) + "/xmls/scene_flat.xml"
@dataclass
class NoiseConfig:
level: float = 1.0
scale_joint_angle: float = 0.03
scale_joint_vel: float = 1.5
scale_gyro: float = 0.2
scale_gravity: float = 0.05
scale_linvel: float = 0.1
@dataclass
class ControlConfig:
# action scale: target angle = actionScale * action + defaultAngle
action_scale = 0.05
@dataclass
class InitState:
# the initial position of the robot in the world frame
pos = [0.0, 0.0, 0.42] #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]
}
@dataclass
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
]
@dataclass
class Normalization:
lin_vel = 2
ang_vel = 0.25
dof_pos = 1
dof_vel = 0.05
@dataclass
class Asset:
body_name = "base"
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",
]
ground = "floor"
@dataclass
class Sensor:
local_linvel = "local_linvel"
gyro = "gyro"
@dataclass
class RewardConfig:
scales: dict[str, float] = field(
default_factory=lambda: {
"termination": -0.0,
"tracking_lin_vel": 1.0,
"tracking_ang_vel": 0.5,
"lin_vel_z": -2.0,
"ang_vel_xy": -0.05,
"orientation": -0.0,
"torques": -0.00001,
"dof_vel": -0.0,
"dof_acc": -2.5e-7,
"base_height": -0.0,
"feet_air_time": 1.0,
"collision": -1.0 * 0,
"feet_stumble": -0.0,
"action_rate": -0.001,
"stand_still": -0.0,
"hip_pos": -1,
"calf_pos": -0.3 * 0,
}
)
tracking_sigma: float = 0.25
max_foot_height: float = 0.1
@registry.envcfg("go2-flat-terrain-walk")
@dataclass
class Go2WalkNpEnvCfg(EnvCfg):
max_episode_seconds: float = 20.0
model_file: str = model_file
noise_config: NoiseConfig = field(default_factory=NoiseConfig)
control_config: ControlConfig = field(default_factory=ControlConfig)
reward_config: RewardConfig = field(default_factory=RewardConfig)
init_state: InitState = field(default_factory=InitState)
commands: Commands = field(default_factory=Commands)
normalization: Normalization = field(default_factory=Normalization)
asset: Asset = field(default_factory=Asset)
sensor: Sensor = field(default_factory=Sensor)
sim_dt: float = 0.01
ctrl_dt: float = 0.01

409
motrix_envs/src/motrix_envs/locomotion/go2/walk_np.py Normal file
View File

@@ -0,0 +1,409 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import gymnasium as gym
import motrixsim as mtx
import numpy as np
from motrix_envs import registry
from motrix_envs.locomotion.go2.cfg import Go2WalkNpEnvCfg
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
_init_dof_vel: np.ndarray
def __init__(self, cfg: Go2WalkNpEnvCfg, num_envs=1):
super().__init__(cfg, num_envs)
self._init_action_space()
self._init_obs_space()
self._body = self._model.get_body(self.cfg.asset.body_name)
self._num_action = self._action_space.shape[0]
self._num_observation = self._observation_space.shape[0]
self._num_dof_pos = self._model.num_dof_pos
self._num_dof_vel = self._model.num_dof_vel
self._init_dof_vel = np.zeros(
(self._num_dof_vel,),
dtype=np.float32,
)
self._init_dof_pos = self._model.compute_init_dof_pos()
self._init_buffer()
def _init_obs_space(self):
model = self.model
num_dof_vel = model.num_dof_vel # linvel + gyro + joint_vel
num_joint_angle = model.num_dof_pos - 7
num_gravity = 3
num_actions = model.num_actuators
num_command = 3
num_obs = num_dof_vel + num_joint_angle + num_gravity + num_actions + num_command
assert num_obs == 48
self._observation_space = gym.spaces.Box(-np.inf, np.inf, (num_obs,), dtype=np.float32)
def _init_action_space(self):
model = self.model
self._action_space = gym.spaces.Box(
np.array(model.actuator_ctrl_limits[0, :]),
np.array(model.actuator_ctrl_limits[1, :]),
(model.num_actuators,),
dtype=np.float32,
)
@property
def action_space(self) -> gym.spaces.Box:
return self._action_space
@property
def observation_space(self) -> gym.spaces.Box:
return self._observation_space
def get_dof_pos(self, data: mtx.SceneModel):
return self._body.get_joint_dof_pos(data)
def get_dof_vel(self, data: mtx.SceneModel):
return self._body.get_joint_dof_vel(data)
def _init_buffer(self):
cfg = self._cfg
assert isinstance(cfg, Go2WalkNpEnvCfg)
# init buffers
self.reset_buf = np.ones(self._num_envs, dtype=np.bool)
self.gravity_vec = np.array([0, 0, -1], dtype=np.float32)
self.commands_scale = np.array(
(
[
cfg.normalization.lin_vel,
cfg.normalization.lin_vel,
cfg.normalization.ang_vel,
]
),
dtype=np.float32,
)
self.default_angles = np.zeros(self._num_action, dtype=np.float32)
self.hip_indices = []
self.calf_indices = []
for i in range(self._model.num_actuators):
for name in cfg.init_state.default_joint_angles.keys():
if name in self._model.actuator_names[i]:
self.default_angles[i] = cfg.init_state.default_joint_angles[name]
if "hip" in self._model.actuator_names[i]:
self.hip_indices.append(i)
if "calf" in self._model.actuator_names[i]:
self.calf_indices.append(i)
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)
self.termination_contact = None
self.foot = []
for name in cfg.asset.terminate_after_contacts_on:
if self.termination_contact is None:
self.termination_contact = np.array([[self._model.get_geom_index(name), self.ground]], dtype=np.uint32)
else:
self.termination_contact = np.append(
self.termination_contact,
np.array(
[[self._model.get_geom_index(name), self.ground]],
dtype=np.uint32,
),
axis=0,
)
for name in cfg.asset.foot_name:
self.foot.append([self._model.get_geom_index(name), self.ground])
self.num_check = self.termination_contact.shape[0]
self.foot = None
for i in self._model.geom_names:
if i is not None and cfg.asset.foot_name in i:
if self.foot is None:
self.foot = np.array([[self._model.get_geom_index(i), self.ground]], dtype=np.uint32)
else:
self.foot = np.append(
self.foot,
np.array(
[[self._model.get_geom_index(i), self.ground]],
dtype=np.uint32,
),
axis=0,
)
self.foot_check_num = self.foot.shape[0]
self.foot_check = self.foot
self.termination_check = self.termination_contact
def apply_action(self, actions, state):
state.info["last_dof_vel"] = self.get_dof_vel(state.data)
state.info["last_actions"] = state.info["current_actions"]
state.info["current_actions"] = actions
state.data.actuator_ctrls = self._compute_target_jq(actions)
return state
def _compute_target_jq(self, actions):
# Compute target position from actions.
target_jq = actions * self.cfg.control_config.action_scale + self.default_angles
return target_jq
def get_local_linvel(self, data: mtx.SceneData) -> np.ndarray:
return self._model.get_sensor_value(self.cfg.sensor.local_linvel, data)
def get_gyro(self, data: mtx.SceneData) -> np.ndarray:
return self._model.get_sensor_value(self.cfg.sensor.gyro, data)
def update_state(self, state):
state = self.update_observation(state)
state = self.update_terminated(state)
state = self.update_reward(state)
return state
def _get_obs(self, data: mtx.SceneData, info: dict) -> np.ndarray:
linear_vel = self.get_local_linvel(data)
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
local_gravity = quat_rotate_inverse(base_quat, self.gravity_vec)
diff = self.get_dof_pos(data) - self.default_angles
noisy_linvel = linear_vel * self.cfg.normalization.lin_vel
noisy_gyro = gyro * self.cfg.normalization.ang_vel
noisy_joint_angle = diff * self.cfg.normalization.dof_pos
noisy_joint_vel = self.get_dof_vel(data) * self.cfg.normalization.dof_vel
command = info["commands"] * self.commands_scale
last_actions = info["current_actions"]
obs = np.hstack(
[
noisy_linvel,
noisy_gyro,
local_gravity,
noisy_joint_angle,
noisy_joint_vel,
last_actions,
command,
]
)
return obs
def update_observation(self, state: NpEnvState):
data = state.data
obs = self._get_obs(data, state.info)
cquerys = self._model.get_contact_query(data)
foot_contact = cquerys.is_colliding(self.foot_check)
state.info["contacts"] = foot_contact.reshape((self._num_envs, self.foot_check_num))
state.info["feet_air_time"] = self.update_feet_air_time(state.info)
return state.replace(obs=obs)
def update_terminated(self, state: NpEnvState) -> NpEnvState:
data = state.data
cquerys = self._model.get_contact_query(data)
termination_check = cquerys.is_colliding(self.termination_check)
termination_check.reshape((self._num_envs, self.num_check))
terminated = termination_check.any(axis=1)
return state.replace(
terminated=terminated,
)
def update_feet_air_time(self, info: dict):
feet_air_time = info["feet_air_time"]
feet_air_time += self.cfg.ctrl_dt
feet_air_time *= ~info["contacts"]
return feet_air_time
def resample_commands(self, num_envs: int):
commands = np.random.uniform(
low=self.cfg.commands.vel_limit[0],
high=self.cfg.commands.vel_limit[1],
size=(num_envs, 3),
)
return commands
def update_reward(self, state: NpEnvState) -> NpEnvState:
data = state.data
terminated = state.terminated
reward_dict = self._get_reward(data, state.info)
rewards = {k: v * self.cfg.reward_config.scales[k] for k, v in reward_dict.items()}
rwd = sum(rewards.values())
rwd = np.clip(rwd, 0.0, 10000.0)
if "termination" in self.cfg.reward_config.scales:
termination = self._reward_termination(terminated) * self.cfg.reward_config.scales["termination"]
rwd += termination
rwd = np.where(terminated, np.array(0.0), rwd)
return state.replace(reward=rwd)
def reset(self, data) -> tuple[np.ndarray, dict]:
num_reset = data.shape[0]
dof_pos = np.tile(self._init_dof_pos, (num_reset, 1))
dof_vel = np.tile(self._init_dof_vel, (num_reset, 1))
data.reset(self._model)
data.set_dof_vel(dof_vel)
data.set_dof_pos(dof_pos, self._model)
self._model.forward_kinematic(data)
info = {
"current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"commands": self.resample_commands(num_reset),
"last_dof_vel": np.zeros((num_reset, self._num_action), dtype=np.float32),
"feet_air_time": np.zeros((num_reset, self.foot_check_num), dtype=np.float32),
"contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
}
obs = self._get_obs(data, info)
return obs, info
def _get_reward(
self,
data: mtx.SceneData,
info: dict,
) -> dict[str, np.ndarray]:
commands = info["commands"]
return {
"lin_vel_z": self._reward_lin_vel_z(data),
"ang_vel_xy": self._reward_ang_vel_xy(data),
"orientation": self._reward_orientation(data),
"torques": self._reward_torques(data),
"dof_vel": self._reward_dof_vel(data),
"dof_acc": self._reward_dof_acc(data, info),
"action_rate": self._reward_action_rate(info),
"tracking_lin_vel": self._reward_tracking_lin_vel(data, commands),
"tracking_ang_vel": self._reward_tracking_ang_vel(data, commands),
"stand_still": self._reward_stand_still(data, commands),
"hip_pos": self._reward_hip_pos(data, commands),
"calf_pos": self._reward_calf_pos(data, commands),
"feet_air_time": self._reward_feet_air_time(commands, info),
}
# ------------ reward functions----------------
def _reward_lin_vel_z(self, data):
# Penalize z axis base linear velocity
return np.square(self.get_local_linvel(data)[:, 2])
def _reward_ang_vel_xy(self, data):
# Penalize xy axes base angular velocity
return np.sum(np.square(self.get_gyro(data)[:, :2]), axis=1)
def _reward_orientation(self, data):
# Penalize non flat base orientation
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = quat_rotate_inverse(base_quat, self.gravity_vec)
return np.sum(np.square(gravity[:, :2]), axis=1)
def _reward_torques(self, data: mtx.SceneData):
# Penalize torques
return np.sum(np.square(data.actuator_ctrls), axis=1)
def _reward_dof_vel(self, data):
# Penalize dof velocities
return np.sum(np.square(self.get_dof_vel(data)), axis=1)
def _reward_dof_acc(self, data, info):
# Penalize dof accelerations
return np.sum(
np.square((info["last_dof_vel"] - self.get_dof_vel(data)) / self.cfg.ctrl_dt),
axis=1,
)
def _reward_action_rate(self, info: dict):
# Penalize changes in actions
action_diff = info["current_actions"] - info["last_actions"]
return np.sum(np.square(action_diff), axis=1)
def _reward_termination(self, done):
# Terminal reward / penalty
return done
def _reward_feet_air_time(self, commands: np.ndarray, info: dict):
# Reward long steps
feet_air_time = info["feet_air_time"]
first_contact = (feet_air_time > 0.0) * info["contacts"]
# reward only on first contact with the ground
rew_airTime = np.sum((feet_air_time - 0.5) * first_contact, axis=1)
# no reward for zero command
rew_airTime *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
return rew_airTime
def _reward_tracking_lin_vel(self, data, commands: np.ndarray):
# Tracking of linear velocity commands (xy axes)
lin_vel_error = np.sum(np.square(commands[:, :2] - self.get_local_linvel(data)[:, :2]), axis=1)
return np.exp(-lin_vel_error / self.cfg.reward_config.tracking_sigma)
def _reward_tracking_ang_vel(self, data, commands: np.ndarray):
# Tracking of angular velocity commands (yaw)
ang_vel_error = np.square(commands[:, 2] - self.get_gyro(data)[:, 2])
return np.exp(-ang_vel_error / self.cfg.reward_config.tracking_sigma)
def _reward_stand_still(self, data, commands: np.ndarray):
# Penalize motion at zero commands
return np.sum(np.abs(self.get_dof_pos(data) - self.default_angles), axis=1) * (
np.linalg.norm(commands, axis=1) < 0.1
)
def _reward_hip_pos(self, data, commands: np.ndarray):
return (0.8 - np.abs(commands[:, 1])) * np.sum(
np.square(self.get_dof_pos(data)[:, self.hip_indices] - self.default_angles[self.hip_indices]),
axis=1,
)
def _reward_calf_pos(self, data, commands: np.ndarray):
return (0.8 - np.abs(commands[:, 1])) * np.sum(
np.square(self.get_dof_pos(data)[:, self.calf_indices] - self.default_angles[self.calf_indices]),
axis=1,
)

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_1.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_2.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_3.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_4.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_1.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_1.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/foot.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_1.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_1.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_0.obj Normal file
View File

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

3
motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_1.obj Normal file
View File

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

279
motrix_envs/src/motrix_envs/locomotion/go2/xmls/go2_mjx.xml Normal file
View File

@@ -0,0 +1,279 @@
<mujoco model="go2">
<compiler angle="radian" autolimits="true"/>
<option iterations="1" ls_iterations="5" timestep="0.004" integrator="Euler">
<flag eulerdamp="disable"/>
</option>
<custom>
<numeric data="30" name="max_contact_points"/>
<numeric data="12" name="max_geom_pairs"/>
</custom>
<default>
<default class="go2">
<geom condim="1"/>
<joint axis="0 1 0" damping="2" armature="0.01"/>
<general forcerange="-24 24" biastype="affine" gainprm="80 0 0" biasprm="0 -80 -1"/>
<default class="abduction">
<joint axis="1 0 0" range="-1.0472 1.0472"/>
<general ctrlrange="-0.9472 0.9472"/>
</default>
<default class="hip">
<joint range="-1.5708 3.4907"/>
<general ctrlrange="-1.4 2.5"/>
</default>
<default class="knee">
<joint range="-2.7227 -0.83776"/>
<general ctrlrange="-2.6227 -0.84776"/>
</default>
<default class="visual">
<geom type="mesh" contype="0" conaffinity="0" group="0"/>
</default>
<default class="go2foot">
<site pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
<geom rgba="0.231373 0.380392 0.705882 1"/>
</default>
<default class="collision">
<geom group="3"/>
<default class="foot">
<geom type="sphere" size="0.0175" pos="-0.002 0 -0.213" solimp="0.9 .95 0.023" condim="3"/>
</default>
</default>
</default>
</default>
<asset>
<material name="metal" rgba=".9 .95 .95 1"/>
<material name="black" rgba="0 0 0 1"/>
<material name="white" rgba="1 1 1 1"/>
<material name="gray" rgba="0.671705 0.692426 0.774270 1"/>
<mesh file="assets/base_0.obj"/>
<mesh file="assets/base_1.obj"/>
<mesh file="assets/base_2.obj"/>
<mesh file="assets/base_3.obj"/>
<mesh file="assets/base_4.obj"/>
<mesh file="assets/hip_0.obj"/>
<mesh file="assets/hip_1.obj"/>
<mesh file="assets/thigh_0.obj"/>
<mesh file="assets/thigh_1.obj"/>
<mesh file="assets/thigh_mirror_0.obj"/>
<mesh file="assets/thigh_mirror_1.obj"/>
<mesh file="assets/calf_0.obj"/>
<mesh file="assets/calf_1.obj"/>
<mesh file="assets/calf_mirror_0.obj"/>
<mesh file="assets/calf_mirror_1.obj"/>
<mesh file="assets/foot.obj"/>
</asset>
<worldbody>
<light name="spotlight" mode="targetbodycom" target="base" pos="3 0 4"/>
<body name="base" pos="0 0 0.445" childclass="go2">
<camera name="track" pos="0.846 -1.465 0.916" xyaxes="0.866 0.500 0.000 -0.171 0.296 0.940" mode="trackcom"/>
<camera name="head_camera" pos="0.32 0 0.04" xyaxes="0 -1 0 0 0 1"/>
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921"
diaginertia="0.107027 0.0980771 0.0244531"/>
<freejoint/>
<!-- <body name="base_0">
<geom mesh="base_0" material="black" class="visual"/>
</body> -->
<geom mesh="base_0" material="black" class="visual"/>
<geom mesh="base_1" material="black" class="visual"/>
<geom mesh="base_2" material="black" class="visual"/>
<geom mesh="base_3" material="white" class="visual"/>
<geom mesh="base_4" material="gray" class="visual"/>
<geom name="base_collision_0" size="0.057 0.04675 0.057" class="collision"/>
<geom name="base_collision_1" size="0.05 0.045" pos="0.285 0 0.01" class="collision"/>
<geom name="base_collision_2" size="0.047" pos="0.293 0 -0.06" class="collision"/>
<site name="imu" pos="-0.02557 0 0.04232"/>
<site name="lidar" pos="0.295 0 -0.075" quat="0.05 0 1 0" rgba="1 0 0 1" size="0.005" group="5"/>
<body name="FL_hip" pos="0.1934 0.0465 0">
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967"/>
<joint name="FL_hip_joint" class="abduction"/>
<geom mesh="hip_0" material="metal" class="visual"/>
<geom mesh="hip_1" material="gray" class="visual"/>
<geom name="fl_hip_0" size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" class="collision"/>
<body name="FL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787"/>
<joint name="FL_thigh_joint" class="hip"/>
<geom mesh="thigh_0" material="metal" class="visual"/>
<geom mesh="thigh_1" material="gray" class="visual"/>
<geom name="fl_thigh_0" size="0.017 0.01225 0.017" pos="-0.015 0 -0.1955" quat="0.707107 0 0.707107 0" class="collision"/>
<body name="FL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
<joint name="FL_calf_joint" class="knee"/>
<geom mesh="calf_0" material="gray" class="visual"/>
<geom mesh="calf_1" material="black" class="visual"/>
<geom name="fl_calf_0" size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" class="collision"/>
<geom name="fl_calf_1" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" class="collision"/>
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
<geom name="FL_foot" class="foot"/>
<site name="FL" class="go2foot"/>
</body>
</body>
</body>
<body name="FR_hip" pos="0.1934 -0.0465 0">
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967"/>
<joint name="FR_hip_joint" class="abduction"/>
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0"/>
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0"/>
<geom name="fr_hip_0" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" class="collision"/>
<body name="FR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787"/>
<joint name="FR_thigh_joint" class="hip"/>
<geom mesh="thigh_mirror_0" material="metal" class="visual"/>
<geom mesh="thigh_mirror_1" material="gray" class="visual"/>
<geom name="fr_thigh_0" size="0.017 0.01225 0.017" pos="-0.015 0 -0.1955" quat="0.707107 0 0.707107 0" class="collision"/>
<body name="FR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
<joint name="FR_calf_joint" class="knee"/>
<geom mesh="calf_mirror_0" material="gray" class="visual"/>
<geom mesh="calf_mirror_1" material="black" class="visual"/>
<geom name="fr_calf_0" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" class="collision"/>
<geom name="fr_calf_1" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" class="collision"/>
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
<geom name="FR_foot" class="foot"/>
<site name="FR" class="go2foot"/>
</body>
</body>
</body>
<body name="RL_hip" pos="-0.1934 0.0465 0">
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967"/>
<joint name="RL_hip_joint" class="abduction"/>
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0"/>
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0"/>
<geom name="rl_hip_0" size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" class="collision"/>
<body name="RL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787"/>
<joint name="RL_thigh_joint" class="hip"/>
<geom mesh="thigh_0" material="metal" class="visual"/>
<geom mesh="thigh_1" material="gray" class="visual"/>
<geom name="rl_thigh_0" size="0.017 0.01225 0.017" pos="-0.015 0 -0.1955" quat="0.707107 0 0.707107 0" class="collision"/>
<body name="RL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
<joint name="RL_calf_joint" class="knee"/>
<geom mesh="calf_0" material="gray" class="visual"/>
<geom mesh="calf_1" material="black" class="visual"/>
<geom name="rl_calf_0" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" class="collision"/>
<geom name="rl_calf_1" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" class="collision"/>
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
<geom name="RL_foot" class="foot"/>
<site name="RL" class="go2foot"/>
</body>
</body>
</body>
<body name="RR_hip" pos="-0.1934 -0.0465 0">
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967"/>
<joint name="RR_hip_joint" class="abduction"/>
<geom mesh="hip_0" material="metal" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1"/>
<geom mesh="hip_1" material="gray" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1"/>
<geom name="rr_hip_0" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" class="collision"/>
<body name="RR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787"/>
<joint name="RR_thigh_joint" class="hip"/>
<geom mesh="thigh_mirror_0" material="metal" class="visual"/>
<geom mesh="thigh_mirror_1" material="gray" class="visual"/>
<geom name="rr_thigh_0" size="0.017 0.01225 0.017" pos="-0.015 0 -0.1955" quat="0.707107 0 0.707107 0" class="collision"/>
<body name="RR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
<joint name="RR_calf_joint" class="knee"/>
<geom mesh="calf_mirror_0" material="gray" class="visual"/>
<geom mesh="calf_mirror_1" material="black" class="visual"/>
<geom name="rr_calf_0" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" class="collision"/>
<geom name="rr_calf_1" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" class="collision"/>
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
<geom name="RR_foot" class="foot"/>
<site name="RR" class="go2foot"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<!-- <motor class="abduction" name="FL_hip" joint="FL_hip_joint"/>
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint"/>
<motor class="knee" name="FL_calf" joint="FL_calf_joint"/>
<motor class="abduction" name="FR_hip" joint="FR_hip_joint"/>
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint"/>
<motor class="knee" name="FR_calf" joint="FR_calf_joint"/>
<motor class="abduction" name="RL_hip" joint="RL_hip_joint"/>
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint"/>
<motor class="knee" name="RL_calf" joint="RL_calf_joint"/>
<motor class="abduction" name="RR_hip" joint="RR_hip_joint"/>
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint"/>
<motor class="knee" name="RR_calf" joint="RR_calf_joint"/> -->
<general class="abduction" name="FL_hip" joint="FL_hip_joint"/>
<general class="hip" name="FL_thigh" joint="FL_thigh_joint"/>
<general class="knee" name="FL_calf" joint="FL_calf_joint"/>
<general class="abduction" name="FR_hip" joint="FR_hip_joint"/>
<general class="hip" name="FR_thigh" joint="FR_thigh_joint"/>
<general class="knee" name="FR_calf" joint="FR_calf_joint"/>
<general class="abduction" name="RL_hip" joint="RL_hip_joint"/>
<general class="hip" name="RL_thigh" joint="RL_thigh_joint"/>
<general class="knee" name="RL_calf" joint="RL_calf_joint"/>
<general class="abduction" name="RR_hip" joint="RR_hip_joint"/>
<general class="hip" name="RR_thigh" joint="RR_thigh_joint"/>
<general class="knee" name="RR_calf" joint="RR_calf_joint"/>
</actuator>
<sensor>
<jointpos joint="FL_hip_joint" name="abduction_front_left_pos"/>
<jointpos joint="FL_thigh_joint" name="hip_front_left_pos"/>
<jointpos joint="FL_calf_joint" name="knee_front_left_pos"/>
<jointpos joint="RL_hip_joint" name="abduction_hind_left_pos"/>
<jointpos joint="RL_thigh_joint" name="hip_hind_left_pos"/>
<jointpos joint="RL_calf_joint" name="knee_hind_left_pos"/>
<jointpos joint="FR_hip_joint" name="abduction_front_right_pos"/>
<jointpos joint="FR_thigh_joint" name="hip_front_right_pos"/>
<jointpos joint="FR_calf_joint" name="knee_front_right_pos"/>
<jointpos joint="RR_hip_joint" name="abduction_hind_right_pos"/>
<jointpos joint="RR_thigh_joint" name="hip_hind_right_pos"/>
<jointpos joint="RR_calf_joint" name="knee_hind_right_pos"/>
<jointvel joint="FL_hip_joint" name="abduction_front_left_vel"/>
<jointvel joint="FL_thigh_joint" name="hip_front_left_vel"/>
<jointvel joint="FL_calf_joint" name="knee_front_left_vel"/>
<jointvel joint="RL_hip_joint" name="abduction_hind_left_vel"/>
<jointvel joint="RL_thigh_joint" name="hip_hind_left_vel"/>
<jointvel joint="RL_calf_joint" name="knee_hind_left_vel"/>
<jointvel joint="FR_hip_joint" name="abduction_front_right_vel"/>
<jointvel joint="FR_thigh_joint" name="hip_front_right_vel"/>
<jointvel joint="FR_calf_joint" name="knee_front_right_vel"/>
<jointvel joint="RR_hip_joint" name="abduction_hind_right_vel"/>
<jointvel joint="RR_thigh_joint" name="hip_hind_right_vel"/>
<jointvel joint="RR_calf_joint" name="knee_hind_right_vel"/>
<gyro site="imu" name="gyro"/>
<accelerometer site="imu" name="accelerometer"/>
<velocimeter site="imu" name="local_linvel"/>
<framezaxis objtype="site" objname="imu" name="upvector"/>
<framequat objtype="site" objname="imu" name="orientation"/>
<framepos objtype="site" objname="imu" name="global_position"/>
<framelinvel objtype="site" objname="imu" name="global_linvel"/>
<frameangvel objtype="site" objname="imu" name="global_angvel"/>
<framelinvel objtype="site" objname="FR" name="FR_global_linvel"/>
<framelinvel objtype="site" objname="FL" name="FL_global_linvel"/>
<framelinvel objtype="site" objname="RR" name="RR_global_linvel"/>
<framelinvel objtype="site" objname="RL" name="RL_global_linvel"/>
<framepos objtype="site" objname="FR" name="FR_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="FL" name="FL_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="RR" name="RR_pos" reftype="site" refname="imu"/>
<framepos objtype="site" objname="RL" name="RL_pos" reftype="site" refname="imu"/>
</sensor>
</mujoco>

39
motrix_envs/src/motrix_envs/locomotion/go2/xmls/scene_flat.xml Normal file
View File

@@ -0,0 +1,39 @@
<mujoco model="go2 feetonly flat terrain scene">
<include file="go2_mjx.xml"/>
<statistic center="0 0 0.1" extent="0.8" meansize="0.04"/>
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global azimuth="120" elevation="-20"/>
<map force="0.01"/>
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2"/>
<quality shadowsize="8192"/>
</visual>
<asset>
<texture type="skybox" builtin="gradient" rgb1="0.4314 0.5294 0.6431" rgb2="0 0 0" width="512" height="512"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.4314 0.5294 0.6431" rgb2="0.8157 0.8549 0.9059"
markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="1 1" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.01" type="plane" material="groundplane" contype="1" conaffinity="0" priority="1"
friction="0.6" condim="3"/>
</worldbody>
<keyframe>
<key name="home"
qpos="
0 0 0.278
1 0 0 0
0.1 0.9 -1.8
-0.1 0.9 -1.8
0.1 0.9 -1.8
-0.1 0.9 -1.8"
ctrl="0.1 0.9 -1.8 -0.1 0.9 -1.8 0.1 0.9 -1.8 -0.1 0.9 -1.8"/>
</keyframe>
</mujoco>

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

@@ -193,7 +193,27 @@ class locomotion:
seed: int = 42
share_policy_value_features: bool = False
max_env_steps: int = 1024 * 60000
max_env_steps: int = 1024 * 60_000
num_envs: int = 2048
# Override PPO configuration
rollouts: int = 24
policy_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
value_hidden_layer_sizes: tuple[int, ...] = (256, 128, 64)
learning_epochs: int = 5
mini_batches: int = 3
learning_rate: float = 3e-4
@rlcfg("go2-flat-terrain-walk")
@dataclass
class Go2WalkPPO(PPOCfg):
"""
Go2 Walk RL config
"""
seed: int = 42
share_policy_value_features: bool = False
max_env_steps: int = 1024 * 60_000
num_envs: int = 2048
# Override PPO configuration