Merge pull request #7 from KraHsu/fix/unitree_example

Align G1 velocity/motion with mjlab and fix doubled log path

Author: KraHsu

examples/unitree/src/genelab_unitree/g1/env_cfg.py

@@ -26,6 +26,46 @@
 _IMU_OFFSET = (0.04525, 0.0, -0.08339)
 _G1_FOOT_LINKS = ("left_ankle_roll_link", "right_ankle_roll_link")
 
+# Per-joint posture standard deviations used by the ``pose`` reward. Smaller std =
+# tighter tolerance. Lower body looser at running speed; arms/wrists looser to allow
+# natural swing. Mirrors mjlab's G1 std table verbatim.
+_G1_POSE_STD_WALKING: dict[str, float] = {
+    # Lower body.
+    r".*hip_pitch.*": 0.3,
+    r".*hip_roll.*": 0.15,
+    r".*hip_yaw.*": 0.15,
+    r".*knee.*": 0.35,
+    r".*ankle_pitch.*": 0.25,
+    r".*ankle_roll.*": 0.1,
+    # Waist.
+    r".*waist_yaw.*": 0.2,
+    r".*waist_roll.*": 0.08,
+    r".*waist_pitch.*": 0.1,
+    # Arms.
+    r".*shoulder_pitch.*": 0.15,
+    r".*shoulder_roll.*": 0.15,
+    r".*shoulder_yaw.*": 0.1,
+    r".*elbow.*": 0.15,
+    r".*wrist.*": 0.3,
+}
+
+_G1_POSE_STD_RUNNING: dict[str, float] = {
+    r".*hip_pitch.*": 0.5,
+    r".*hip_roll.*": 0.2,
+    r".*hip_yaw.*": 0.2,
+    r".*knee.*": 0.6,
+    r".*ankle_pitch.*": 0.35,
+    r".*ankle_roll.*": 0.15,
+    r".*waist_yaw.*": 0.3,
+    r".*waist_roll.*": 0.08,
+    r".*waist_pitch.*": 0.2,
+    r".*shoulder_pitch.*": 0.5,
+    r".*shoulder_roll.*": 0.2,
+    r".*shoulder_yaw.*": 0.15,
+    r".*elbow.*": 0.35,
+    r".*wrist.*": 0.3,
+}
+
 
 def _obs_terms() -> dict[str, ObservationTermCfg]:
     return {
@@ -128,8 +168,8 @@ def unitree_g1_velocity_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
                 heading_control_stiffness=0.5,
                 ranges=UniformVelocityCommandCfg.Ranges(
                     lin_vel_x=(-1.0, 1.0),
-                    lin_vel_y=(-0.5, 0.5),
-                    ang_vel_z=(-0.7, 0.7),
+                    lin_vel_y=(-1.0, 1.0),
+                    ang_vel_z=(-0.5, 0.5),
                     heading=(-math.pi, math.pi),
                 ),
             )
@@ -145,12 +185,26 @@ def unitree_g1_velocity_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
                 weight=2.0,
                 params={"command_name": "twist", "std": math.sqrt(0.5)},
             ),
-            "upright": RewardTermCfg(func=mdp.flat_orientation_l2, weight=-1.0),
+            "upright": RewardTermCfg(
+                func=mdp.upright_exp,
+                weight=1.0,
+                params={"std": math.sqrt(0.2)},
+            ),
+            "pose": RewardTermCfg(
+                func=mdp.variable_posture,
+                weight=1.0,
+                params={
+                    "command_name": "twist",
+                    "walking_threshold": 0.05,
+                    "running_threshold": 1.5,
+                    "default_std": 0.3,
+                    "std_standing": {".*": 0.05},
+                    "std_walking": _G1_POSE_STD_WALKING,
+                    "std_running": _G1_POSE_STD_RUNNING,
+                },
+            ),
             "action_rate": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.1),
             "dof_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-1.0),
-            "feet_air_time": RewardTermCfg(
-                func=mdp.feet_air_time, weight=0.5, params={"threshold": 0.4}
-            ),
         },
         terminations_cfg={
             "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
@@ -163,7 +217,12 @@ def unitree_g1_velocity_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
                 mode="reset",
                 func=mdp.reset_root_state_uniform,
                 params={
-                    "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (0.01, 0.05)},
+                    "pose_range": {
+                        "x": (-0.5, 0.5),
+                        "y": (-0.5, 0.5),
+                        "z": (0.01, 0.05),
+                        "yaw": (-math.pi, math.pi),
+                    },
                     "velocity_range": {},
                 },
             ),
@@ -173,8 +232,17 @@ def unitree_g1_velocity_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
     if not play:
         cfg.events_cfg["push_robot"] = EventTermCfg(
             mode="interval",
-            interval_range_s=(10.0, 15.0),
+            interval_range_s=(1.0, 3.0),
             func=mdp.push_by_setting_velocity,
-            params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
+            params={
+                "velocity_range": {
+                    "x": (-0.5, 0.5),
+                    "y": (-0.5, 0.5),
+                    "z": (-0.4, 0.4),
+                    "roll": (-0.52, 0.52),
+                    "pitch": (-0.52, 0.52),
+                    "yaw": (-0.78, 0.78),
+                },
+            },
         )
     return cfg

examples/unitree/src/genelab_unitree/g1/ppo_cfg.py

@@ -7,7 +7,7 @@ def unitree_g1_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
     return RslRlOnPolicyRunnerCfg(
         seed=42,
         num_steps_per_env=24,
-        max_iterations=3000,
+        max_iterations=30_000,
         save_interval=50,
         experiment_name="g1_velocity_flat",
         logger="tensorboard",

examples/unitree/src/genelab_unitree/g1/tracking_env_cfg.py

@@ -227,6 +227,15 @@ def unitree_g1_tracking_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
             mode="interval",
             interval_range_s=(1.0, 3.0),
             func=mdp.push_by_setting_velocity,
-            params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
+            params={
+                "velocity_range": {
+                    "x": (-0.5, 0.5),
+                    "y": (-0.5, 0.5),
+                    "z": (-0.4, 0.4),
+                    "roll": (-0.52, 0.52),
+                    "pitch": (-0.52, 0.52),
+                    "yaw": (-0.78, 0.78),
+                },
+            },
         )
     return cfg

src/genelab/cli/__init__.py

@@ -388,7 +388,7 @@ def _dispatch_train(task: _RunnableTask, runner_args: dict[str, str]) -> None:
         num_envs=int(num_envs_raw) if num_envs_raw is not None else None,
         max_iterations=int(max_iter_raw) if max_iter_raw is not None else None,
         seed=int(seed_raw) if seed_raw is not None else None,
-        log_root=Path(log_dir_raw) if log_dir_raw is not None else None,
+        log_dir=Path(log_dir_raw) if log_dir_raw is not None else None,
     )
 
 

src/genelab/mdp/__init__.py

@@ -46,6 +46,8 @@
     motion_relative_body_position_error_exp,
     track_angular_velocity_z_exp,
     track_linear_velocity_xy_exp,
+    upright_exp,
+    variable_posture,
 )
 from genelab.mdp.terminations import (
     bad_anchor_ori,
@@ -105,4 +107,6 @@
     "time_out",
     "track_angular_velocity_z_exp",
     "track_linear_velocity_xy_exp",
+    "upright_exp",
+    "variable_posture",
 ]

src/genelab/mdp/events.py

@@ -4,6 +4,8 @@
 
 import torch
 
+from genelab.utils.math import quat_from_euler_xyz, quat_mul
+
 if TYPE_CHECKING:
     from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv
 
@@ -14,7 +16,13 @@ def reset_root_state_uniform(
     pose_range: dict[str, tuple[float, float]] | None = None,
     velocity_range: dict[str, tuple[float, float]] | None = None,
 ) -> None:
-    """Randomize the floating base pose / velocity within the given ranges."""
+    """Randomize the floating base pose / velocity within the given ranges.
+
+    ``pose_range`` accepts ``x``, ``y``, ``z`` (position offsets, m) and ``roll``, ``pitch``,
+    ``yaw`` (rotation offsets, rad, layered on top of the configured ``init_quat``).
+    ``velocity_range`` accepts ``x``, ``y``, ``z`` (linear, m/s) and ``roll``, ``pitch``,
+    ``yaw`` (angular, rad/s).
+    """
     if env_ids is None or env_ids.numel() == 0:
         return
     pose_range = pose_range or {}
@@ -31,12 +39,34 @@ def reset_root_state_uniform(
             set_pos(init_pos, envs_idx=env_ids)
         except TypeError:
             set_pos(init_pos)
+    if any(axis in pose_range for axis in ("roll", "pitch", "yaw")):
+        base_quat = torch.tensor(env.cfg.robot.init_quat, device=env.device).expand(n, -1).clone()
+        zeros = torch.zeros(n, device=env.device)
+        rpy = []
+        for axis in ("roll", "pitch", "yaw"):
+            if axis in pose_range:
+                lo, hi = pose_range[axis]
+                rpy.append(torch.empty(n, device=env.device).uniform_(lo, hi))
+            else:
+                rpy.append(zeros)
+        rand_quat = quat_from_euler_xyz(rpy[0], rpy[1], rpy[2])
+        new_quat = quat_mul(base_quat, rand_quat)
+        set_quat = getattr(env.robot, "set_quat", None)
+        if set_quat is not None:
+            try:
+                set_quat(new_quat, envs_idx=env_ids)
+            except TypeError:
+                set_quat(new_quat)
     vel = torch.zeros(n, 3, device=env.device)
     ang = torch.zeros(n, 3, device=env.device)
     for axis, idx in (("x", 0), ("y", 1), ("z", 2)):
         if axis in velocity_range:
             lo, hi = velocity_range[axis]
             vel[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi)
+    for axis, idx in (("roll", 0), ("pitch", 1), ("yaw", 2)):
+        if axis in velocity_range:
+            lo, hi = velocity_range[axis]
+            ang[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi)
     set_vel = getattr(env.robot, "set_vel", None)
     set_ang = getattr(env.robot, "set_ang", None)
     if set_vel is not None:
@@ -86,18 +116,30 @@ def push_by_setting_velocity(
     env_ids: torch.Tensor | None,
     velocity_range: dict[str, tuple[float, float]] | None = None,
 ) -> None:
+    """Push selected envs by overwriting linear (x/y/z) and angular (roll/pitch/yaw) base velocity."""
     if env_ids is None or env_ids.numel() == 0:
         return
     velocity_range = velocity_range or {}
     n = int(env_ids.numel())
     vel = torch.zeros(n, 3, device=env.device)
+    ang = torch.zeros(n, 3, device=env.device)
     for axis, idx in (("x", 0), ("y", 1), ("z", 2)):
         if axis in velocity_range:
             lo, hi = velocity_range[axis]
             vel[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi)
+    for axis, idx in (("roll", 0), ("pitch", 1), ("yaw", 2)):
+        if axis in velocity_range:
+            lo, hi = velocity_range[axis]
+            ang[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi)
     set_vel = getattr(env.robot, "set_vel", None)
+    set_ang = getattr(env.robot, "set_ang", None)
     if set_vel is not None:
         try:
             set_vel(vel, envs_idx=env_ids)
         except TypeError:
             set_vel(vel)
+    if set_ang is not None:
+        try:
+            set_ang(ang, envs_idx=env_ids)
+        except TypeError:
+            set_ang(ang)

src/genelab/mdp/rewards.py

@@ -1,9 +1,11 @@
 """Reusable reward term functions for locomotion tasks."""
 
+import re
 from typing import TYPE_CHECKING, cast
 
 import torch
 
+from genelab.managers.reward_manager import RewardTermCfg
 from genelab.mdp.commands.motion_command import MotionCommand
 from genelab.utils.math import quat_error_magnitude
 
@@ -45,6 +47,93 @@ def flat_orientation_l2(env: "ManagerBasedRlEnv") -> torch.Tensor:
     return torch.sum(env.robot_state.projected_gravity_b[:, :2] ** 2, dim=-1)
 
 
+def upright_exp(env: "ManagerBasedRlEnv", std: float = 0.45) -> torch.Tensor:
+    """``exp(-||projected_gravity_xy||^2 / std^2)`` — positive reward for an upright base.
+
+    Port of mjlab's ``upright`` reward (flat-ground variant). Saturates near zero tilt
+    instead of growing unbounded like an L2 penalty, which matches the policy gradients
+    the reference implementation relies on.
+    """
+    xy_squared = torch.sum(env.robot_state.projected_gravity_b[:, :2] ** 2, dim=-1)
+    return torch.exp(-xy_squared / (std * std))
+
+
+class variable_posture:
+    """Reward for staying near the default pose, with per-joint, speed-dependent tolerance.
+
+    Port of ``mjlab.tasks.velocity.mdp.rewards.variable_posture``. Std dicts map joint name
+    regex → std value; per joint the *last* matching pattern wins (same convention as
+    ``RobotEntityCfg.default_joint_pos``). Joints with no match keep the supplied ``default``.
+
+    At each step the active std vector is chosen per env from the command magnitude:
+
+    * ``total_speed = ||lin_vel_xy|| + |ang_vel_z|``
+    * standing if ``total_speed < walking_threshold``
+    * walking if ``walking_threshold <= total_speed < running_threshold``
+    * running otherwise
+
+    Reward: ``exp(-mean((joint_pos - default)^2 / std^2))``.
+    """
+
+    def __init__(self, cfg: RewardTermCfg, env: "ManagerBasedRlEnv") -> None:
+        self._env = env
+        params = cfg.params
+        default = float(params.get("default_std", 1.0))
+        std_standing = params.get("std_standing", {})
+        std_walking = params.get("std_walking", {})
+        std_running = params.get("std_running", {})
+        self._std_standing = self._build_std(std_standing, default)
+        self._std_walking = self._build_std(std_walking, default)
+        self._std_running = self._build_std(std_running, default)
+
+    def _build_std(self, mapping: dict[str, float], default: float) -> torch.Tensor:
+        out = torch.full((len(self._env.joint_names),), default, device=self._env.device)
+        for pattern, value in mapping.items():
+            try:
+                regex = re.compile(pattern)
+            except re.error:
+                regex = re.compile(re.escape(pattern))
+            for i, name in enumerate(self._env.joint_names):
+                if regex.fullmatch(name) or regex.search(name):
+                    out[i] = float(value)
+        return out
+
+    def __call__(
+        self,
+        env: "ManagerBasedRlEnv",
+        command_name: str,
+        std_standing: dict[str, float] | None = None,
+        std_walking: dict[str, float] | None = None,
+        std_running: dict[str, float] | None = None,
+        default_std: float = 1.0,
+        walking_threshold: float = 0.05,
+        running_threshold: float = 1.5,
+    ) -> torch.Tensor:
+        del std_standing, std_walking, std_running, default_std  # consumed at __init__
+
+        command = env.command_manager.get_command(command_name)
+        linear_speed = torch.norm(command[:, :2], dim=-1)
+        angular_speed = torch.abs(command[:, 2])
+        total_speed = linear_speed + angular_speed
+
+        standing_mask = (total_speed < walking_threshold).float().unsqueeze(-1)
+        walking_mask = (
+            ((total_speed >= walking_threshold) & (total_speed < running_threshold))
+            .float()
+            .unsqueeze(-1)
+        )
+        running_mask = (total_speed >= running_threshold).float().unsqueeze(-1)
+
+        std = (
+            self._std_standing * standing_mask
+            + self._std_walking * walking_mask
+            + self._std_running * running_mask
+        )
+
+        error = env.robot_state.joint_pos - env.default_joint_pos
+        return torch.exp(-torch.mean((error * error) / (std * std), dim=-1))
+
+
 def joint_pos_limits(env: "ManagerBasedRlEnv") -> torch.Tensor:
     """L2 of joint-position excursion past ±π (cheap stand-in for true limit penalty)."""
     excess = (env.robot_state.joint_pos.abs() - 3.14).clamp(min=0.0)

src/genelab/rl/runner.py

@@ -97,9 +97,15 @@ def train_task(
     max_iterations: int | None = None,
     seed: int | None = None,
     log_root: Path | None = None,
+    log_dir: Path | None = None,
     resume_from: Path | None = None,
 ) -> Path:
-    """Train ``task_id`` with PPO. Returns the log directory."""
+    """Train ``task_id`` with PPO. Returns the log directory.
+
+    ``log_dir`` (final, pre-resolved) takes precedence over ``log_root`` (parent under
+    which a ``<experiment>/<timestamp>`` directory is created). The torchrun relaunch
+    path pre-resolves the log dir so every rank lands in the same folder.
+    """
     ensure_project_cache()
     env_cfg = _resolve_env_cfg(task_id, play=False)
     if num_envs is not None:
@@ -115,8 +121,9 @@ def train_task(
 
     from rsl_rl.runners import OnPolicyRunner
 
-    log_root = log_root or Path("logs") / "rsl_rl"
-    log_dir = resolve_log_dir(log_root, agent_cfg.experiment_name, agent_cfg.run_name)
+    if log_dir is None:
+        log_root = log_root or Path("logs") / "rsl_rl"
+        log_dir = resolve_log_dir(log_root, agent_cfg.experiment_name, agent_cfg.run_name)
     if is_main_process():
         _save_run_params(log_dir, env_cfg, agent_cfg)