diff --git a/motrix_envs/src/motrix_envs/locomotion/__init__.py b/motrix_envs/src/motrix_envs/locomotion/__init__.py index b56709a..c829201 100644 --- a/motrix_envs/src/motrix_envs/locomotion/__init__.py +++ b/motrix_envs/src/motrix_envs/locomotion/__init__.py @@ -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 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/__init__.py b/motrix_envs/src/motrix_envs/locomotion/go2/__init__.py new file mode 100644 index 0000000..dbc0c8e --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/__init__.py @@ -0,0 +1,16 @@ +# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from . import walk_np # noqa: F401 register envs diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/cfg.py b/motrix_envs/src/motrix_envs/locomotion/go2/cfg.py new file mode 100644 index 0000000..be6f19d --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/cfg.py @@ -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 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/go2/walk_np.py new file mode 100644 index 0000000..ef66259 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/walk_np.py @@ -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, + ) diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_0.obj new file mode 100644 index 0000000..5fe9df9 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cdc53719002a5ee737d867904e534a6d6fa930b616bdcc1b3131e577d4e46c0 +size 1330465 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_1.obj new file mode 100644 index 0000000..794801d --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:39d2f4724ff64dd5a0640c308fd98902a2946660a8d033ec8dbf368fd5af0dd6 +size 811791 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_2.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_2.obj new file mode 100644 index 0000000..4f9b38d --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_2.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:29cae463f95d678e4ede863c36c9b19b0d356b8340412add10670d70bdcc3095 +size 294090 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_3.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_3.obj new file mode 100644 index 0000000..a15aa64 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_3.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f51ba43bf300e68957a0c1aa83eaf3e4926f4ac7e64a0d6c33ae5d21fcbcf5f6 +size 379093 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_4.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_4.obj new file mode 100644 index 0000000..b1f8c0d --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_4.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49ddec0129546e5dd2075d5beacf7f207baa2860aeff7626818e64d2efbf3821 +size 7776234 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_0.obj new file mode 100644 index 0000000..3f35628 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4a818cd563c5b7e2a271504442d05fd777d38cbaa028d4a10ad29c9f87bb355e +size 877174 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_1.obj new file mode 100644 index 0000000..053c95f --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:51950c3a153ae3208592df86ea11c9d6756f8f645ebca5caf6910c61aeb03817 +size 326865 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_0.obj new file mode 100644 index 0000000..ba259c3 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6dc0406207a1c08a5953fd5962aca1a270fd06593d30d100205e05c1b32d4af2 +size 876538 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_1.obj new file mode 100644 index 0000000..6a8527e --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f84f5af90b914633165beb3a5493ee499169144c5797dcf04c6eb5842f74b31a +size 327261 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/foot.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/foot.obj new file mode 100644 index 0000000..51a8bf8 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/foot.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:df9e78a7c0110d02557439010f2a5f11f0ec4fd907b32704145dbd6aa7affc0b +size 1144729 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_0.obj new file mode 100644 index 0000000..cc11c34 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc29932ab4d251ce3f72389ec4c1f757e44eee138b87e772b648c2bc76b506cd +size 2863864 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_1.obj new file mode 100644 index 0000000..ce4ab76 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3fd5bd9b2ef6abe622c0739d3c67eaa1679b19b47132951972c4892a779ea7cf +size 2767918 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_0.obj new file mode 100644 index 0000000..89d373e --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb3edfe6a7b04f4f4dff1ce7b05d426ddce941e276cc6372fb542055893c71ec +size 2841868 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_1.obj new file mode 100644 index 0000000..7a231d1 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3a631822f9f7d2114b1b79dcdff5abf95103630abb7d91d1f3b8b18117132ed2 +size 1491320 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_0.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_0.obj new file mode 100644 index 0000000..12995f2 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_0.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07abaedce761f1dd48647384ac14ba1959276bc78d5a2a62051e8f1c6093933b +size 2817674 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_1.obj b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_1.obj new file mode 100644 index 0000000..3417868 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:db984a3e983eabbc9b38183427cca9ab0bb3bdb049f9eab030e06edda36579d6 +size 1482261 diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/go2_mjx.xml b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/go2_mjx.xml new file mode 100644 index 0000000..c0035a2 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/go2_mjx.xml @@ -0,0 +1,279 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/motrix_envs/src/motrix_envs/locomotion/go2/xmls/scene_flat.xml b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/scene_flat.xml new file mode 100644 index 0000000..719d60e --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go2/xmls/scene_flat.xml @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/motrix_rl/src/motrix_rl/cfgs.py b/motrix_rl/src/motrix_rl/cfgs.py index 0a5a660..1e002c7 100644 --- a/motrix_rl/src/motrix_rl/cfgs.py +++ b/motrix_rl/src/motrix_rl/cfgs.py @@ -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