add go2 training

Author: jyf23

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

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

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