From dbfa9e31fa41d968311f940d288794775235a1d0 Mon Sep 17 00:00:00 2001
From: jyf23 <1305116641@qq.com>
Date: Fri, 2 Jan 2026 21:10:34 +0800
Subject: [PATCH] add go2 training

---
 .../src/motrix_envs/locomotion/__init__.py    |   2 +-
 .../motrix_envs/locomotion/go2/__init__.py    |  16 +
 .../src/motrix_envs/locomotion/go2/cfg.py     | 139 ++++++
 .../src/motrix_envs/locomotion/go2/walk_np.py | 409 ++++++++++++++++++
 .../locomotion/go2/xmls/assets/base_0.obj     |   3 +
 .../locomotion/go2/xmls/assets/base_1.obj     |   3 +
 .../locomotion/go2/xmls/assets/base_2.obj     |   3 +
 .../locomotion/go2/xmls/assets/base_3.obj     |   3 +
 .../locomotion/go2/xmls/assets/base_4.obj     |   3 +
 .../locomotion/go2/xmls/assets/calf_0.obj     |   3 +
 .../locomotion/go2/xmls/assets/calf_1.obj     |   3 +
 .../go2/xmls/assets/calf_mirror_0.obj         |   3 +
 .../go2/xmls/assets/calf_mirror_1.obj         |   3 +
 .../locomotion/go2/xmls/assets/foot.obj       |   3 +
 .../locomotion/go2/xmls/assets/hip_0.obj      |   3 +
 .../locomotion/go2/xmls/assets/hip_1.obj      |   3 +
 .../locomotion/go2/xmls/assets/thigh_0.obj    |   3 +
 .../locomotion/go2/xmls/assets/thigh_1.obj    |   3 +
 .../go2/xmls/assets/thigh_mirror_0.obj        |   3 +
 .../go2/xmls/assets/thigh_mirror_1.obj        |   3 +
 .../locomotion/go2/xmls/go2_mjx.xml           | 279 ++++++++++++
 .../locomotion/go2/xmls/scene_flat.xml        |  39 ++
 motrix_rl/src/motrix_rl/cfgs.py               |  22 +-
 23 files changed, 952 insertions(+), 2 deletions(-)
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/__init__.py
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/cfg.py
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/walk_np.py
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_2.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_3.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/base_4.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/calf_mirror_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/foot.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/hip_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_0.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/assets/thigh_mirror_1.obj
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/go2_mjx.xml
 create mode 100644 motrix_envs/src/motrix_envs/locomotion/go2/xmls/scene_flat.xml

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 @@
+<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>
\ 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 @@
+<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>
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