Add energy rewards and MujocoStyleActuatorCfg (closes S4)

CHANGELOG.md

@@ -41,6 +41,19 @@ All notable changes to GeneLab are recorded here.
   integration test that exercises all four new sensor wrappers via
   `gs.Scene + add_sensor + build + step`, validating the wrappers'
   `pre_build_genesis` / `_compute_data` round-trip end-to-end.
+- **Energy rewards.** Added `genelab.mdp.rewards.{kinetic_energy_l2,
+  potential_energy, energy_budget}` wrapping Genesis 1.0's
+  `gs_handle.{get_kinetic_energy, get_potential_energy, get_total_energy}`.
+  `kinetic_energy_l2` and `energy_budget` are non-negative (squared);
+  `potential_energy` is signed.
+- **MJCF-style actuator cfg.** Added `MujocoStyleActuatorCfg`, a thin
+  cfg over `IdealPDActuator` that accepts Mujoco's `(gear, bias_prm,
+  dyntype="none", gaintype="fixed", biastype="affine")` shape and
+  translates it into `(stiffness, damping, effort_limit)`. Mirrors what
+  Genesis's own MJCF parser does for `<actuator>` elements at load
+  time, so robots whose actuator parameters live in Python config —
+  rather than in an MJCF file — can use the same named knobs without
+  hand-translating to PD gains.
 
 ### Changed
 

docs/concepts/actuators.zh.md

@@ -16,6 +16,9 @@
 | `DCMotorActuator` | 在 ideal PD 基础上加入电机限制和饱和行为。 |
 | `MlpResidualActuator` | 在 `DCMotorActuator` 基础上叠加 TorchScript 残差力矩模型。 |
 
+`MujocoStyleActuatorCfg` 是 `IdealPDActuator` 的便捷配置 —— 详见下文
+[MJCF 风格的 actuator 配置](#mjcf-style-actuator-configuration)。
+
 执行器按配置的关节名或表达式匹配关节组，并向 action term 暴露维度和控制逻辑。
 
 ## 使用 `MlpResidualActuatorCfg`
@@ -54,6 +57,29 @@ TorchScript 模块接收最后一维为 `[target_pos - joint_pos, joint_vel]` 
 genelab play GeneLab-MlpResidual-Actuator-Showcase-v0 --steps 5
 ```
 
+## MJCF 风格的 actuator 配置 { #mjcf-style-actuator-configuration }
+
+`MujocoStyleActuatorCfg` 把 Mujoco general actuator 翻译成 `IdealPDActuator`。当
+机器人在 MJCF 里以 `dyntype=none`、`gaintype=fixed`、`biastype=affine` 形式声明
+actuator，而手写等价 `stiffness` / `damping` 会让原始 source-of-truth 不清晰时使用它。
+
+```python
+from genelab.actuator import MujocoStyleActuatorCfg
+
+robot_cfg.actuators["arm"] = MujocoStyleActuatorCfg(
+    target_names_expr=("joint[1-7]",),
+    gear=1.0,
+    bias_prm=(0.0, -200.0, -5.0),  # 常数项、qpos 项、qvel 项
+    effort_limit=80.0,
+)
+# 等价于 IdealPDActuatorCfg(stiffness=200.0, damping=5.0, effort_limit=80.0)。
+```
+
+任何其它 dyntype / gaintype / biastype 组合都被拒绝（Genesis 自带的 MJCF parser 也对
+同一组合记录 warning），且 `bias_prm[0]` 必须为零 —— 因为 `IdealPDActuator` 没有常数
+偏置力矩项。最终的 `stiffness` / `damping` 计算为 `-gear * bias_prm[1]` 和
+`-gear * bias_prm[2]`。
+
 ## 设计建议
 
 执行器分组应贴合机器人机构。若手臂、手、底座关节需要不同增益、限制或 action scale，不要放进一个巨大的 actuator。

src/genelab/actuator/__init__.py

@@ -13,6 +13,7 @@
 from genelab.actuator.ideal_pd import IdealPDActuator, IdealPDActuatorCfg
 from genelab.actuator.implicit_pd import ImplicitPDActuator, ImplicitPDActuatorCfg
 from genelab.actuator.mlp_residual import MlpResidualActuator, MlpResidualActuatorCfg
+from genelab.actuator.mujoco_style import MujocoStyleActuatorCfg
 
 __all__ = [
     "ActuatorBase",
@@ -25,4 +26,5 @@
     "ImplicitPDActuatorCfg",
     "MlpResidualActuator",
     "MlpResidualActuatorCfg",
+    "MujocoStyleActuatorCfg",
 ]

src/genelab/actuator/mujoco_style.py

@@ -0,0 +1,100 @@
+"""Mujoco-style actuator cfg — convenience translator over :class:`IdealPDActuator`.
+
+Mujoco's general actuator transfer function is
+
+.. code-block:: text
+
+    force = gear * (gain_prm[0] * ctrl
+                    - bias_prm[0] - bias_prm[1] * qpos - bias_prm[2] * qvel)
+
+For the most common MJCF actuator shape — ``dyntype=NONE``, ``gaintype=FIXED``,
+``biastype=AFFINE``, ``gain_prm[0] = 1`` — this reduces to a PD controller:
+
+.. code-block:: text
+
+    kp = -gear * bias_prm[1]
+    kv = -gear * bias_prm[2]
+
+with the controller's ``ctrl`` interpreted as a position target.
+:class:`MujocoStyleActuatorCfg` accepts those Mujoco parameters and translates
+them into the underlying :class:`~genelab.actuator.IdealPDActuator` primitives
+so MJCF-defined actuators can be brought into GeneLab without translating the
+``(kp, kv, force_range)`` triple by hand.
+
+Genesis 1.0 has no separate runtime "general actuator" class — the MJCF parser
+in ``genesis/utils/mjcf.py`` already translates these same Mujoco attributes
+into Genesis's ``set_dofs_kp`` / ``set_dofs_kv`` / ``set_dofs_force_range``
+setters at load time. This cfg is the GeneLab-side mirror: it exposes the same
+named knobs for tasks that configure actuators in Python rather than loading
+them from an MJCF.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+from genelab.actuator.ideal_pd import IdealPDActuator, IdealPDActuatorCfg
+
+
+@dataclass
+class MujocoStyleActuatorCfg(IdealPDActuatorCfg):
+    """Cfg for actuators specified in Mujoco's ``(gear, gain_prm, bias_prm)`` form.
+
+    Only the canonical ``dyntype=none / gaintype=fixed / biastype=affine`` shape is
+    supported — every other dyntype / gaintype / biastype combination raises a
+    ``ValueError`` at construction time (Genesis itself logs a warning and falls
+    back to the supported shape when an MJCF declares them, so rejecting at the
+    Python cfg layer mirrors that behaviour up front).
+
+    ``bias_prm[0]`` (the constant bias term) must be zero — ``IdealPDActuator``
+    has no constant-offset force term, so a non-zero constant cannot be
+    represented exactly. ``stiffness`` / ``damping`` are computed from ``gear``
+    and ``bias_prm``; setting them directly on this cfg raises.
+    """
+
+    # Mujoco's transfer-function knobs.
+    gear: float = 1.0
+    bias_prm: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    # Discrete-choice fields kept on the cfg for documentation; the supported
+    # combination is the only one allowed today. Drop the ``Literal`` constraint
+    # if a future revision needs to expand the supported set.
+    dyntype: Literal["none"] = "none"
+    gaintype: Literal["fixed"] = "fixed"
+    biastype: Literal["affine"] = "affine"
+
+    def __post_init__(self) -> None:
+        if self.dyntype != "none":
+            raise ValueError(
+                f"MujocoStyleActuatorCfg: dyntype={self.dyntype!r} not supported; "
+                "only 'none' (no internal actuator dynamics) is implementable on "
+                "the existing IdealPDActuator primitives."
+            )
+        if self.gaintype != "fixed":
+            raise ValueError(
+                f"MujocoStyleActuatorCfg: gaintype={self.gaintype!r} not supported; "
+                "only 'fixed' (gain_prm[0]=1) is implementable on IdealPDActuator."
+            )
+        if self.biastype != "affine":
+            raise ValueError(
+                f"MujocoStyleActuatorCfg: biastype={self.biastype!r} not supported; "
+                "only 'affine' (kp/kv terms via bias_prm[1] / bias_prm[2]) maps to PD."
+            )
+        bias_const, bias_kp_term, bias_kv_term = self.bias_prm
+        if bias_const != 0.0:
+            raise ValueError(
+                f"MujocoStyleActuatorCfg: bias_prm[0]={bias_const} (constant bias) "
+                "is not representable by IdealPDActuator — there is no constant-offset "
+                "force term. Set bias_prm[0]=0."
+            )
+        if self.stiffness != 0.0 or self.damping != 0.0:
+            raise ValueError(
+                "MujocoStyleActuatorCfg: stiffness / damping are computed from "
+                "gear * bias_prm; don't set them directly on this cfg."
+            )
+        # PD gains derived from the Mujoco bias parameters; the sign flip turns
+        # the +bias_prm[1]*qpos / +bias_prm[2]*qvel terms (which Mujoco *subtracts*
+        # from the control output) into the +kp*(target-qpos) / -kv*qvel terms
+        # IdealPDActuator computes.
+        self.stiffness = -self.gear * bias_kp_term
+        self.damping = -self.gear * bias_kv_term
+        if self.class_type is None:
+            self.class_type = IdealPDActuator

src/genelab/mdp/rewards/__init__.py

@@ -29,6 +29,11 @@
     upright_exp,
     variable_posture,
 )
+from genelab.mdp.rewards.energy import (
+    energy_budget,
+    kinetic_energy_l2,
+    potential_energy,
+)
 from genelab.mdp.rewards.tactile import (
     contact_count,
     contact_intensity_l2,
@@ -54,6 +59,7 @@
     "body_angular_velocity_penalty",
     "contact_count",
     "contact_intensity_l2",
+    "energy_budget",
     "feet_air_time",
     "feet_clearance",
     "feet_slip",
@@ -62,13 +68,15 @@
     "joint_acc_l2",
     "joint_pos_limits",
     "joint_vel_limits",
+    "kinetic_energy_l2",
     "lin_vel_z_l2",
     "motion_global_anchor_orientation_error_exp",
     "motion_global_anchor_position_error_exp",
     "motion_global_body_angular_velocity_error_exp",
     "motion_global_body_linear_velocity_error_exp",
     "motion_relative_body_orientation_error_exp",
     "motion_relative_body_position_error_exp",
+    "potential_energy",
     "self_collision_cost",
     "slip_penalty",
     "soft_landing",

src/genelab/mdp/rewards/energy.py

@@ -0,0 +1,75 @@
+"""Energy-budget reward primitives.
+
+Wraps Genesis 1.0's per-entity energy accessors —
+``Articulation.gs_handle.get_kinetic_energy()`` /
+``get_potential_energy()`` / ``get_total_energy()`` — into ``(num_envs,)``
+reward terms. Useful as soft regularizers in sim-to-real curricula and
+as hard-constraints on motion smoothness.
+
+Each function takes the same ``asset_cfg`` selector pattern as the rest of
+``genelab.mdp.rewards``: ``None`` resolves to the primary ``"robot"`` entity,
+a populated :class:`SceneEntityCfg` routes to its named asset.
+
+Sign conventions:
+
+* :func:`kinetic_energy_l2` is non-negative (squared). Pair with a negative
+  weight to penalize high-energy motion (motion-smoothness regularizer).
+* :func:`potential_energy` is signed (Genesis returns negative values for
+  centres-of-mass at world-frame height < 0); pair with a positive or
+  negative weight depending on whether the task wants the body raised or
+  lowered. The default ``base_height_l2`` is usually a better fit for
+  base-height tracking; this primitive is for tasks that want the *full
+  articulation*'s gravitational potential.
+* :func:`energy_budget` is non-negative (squared deviation from a target
+  total). Pair with a negative weight.
+"""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from genelab.mdp._helpers import asset_handle as _asset_handle
+
+if TYPE_CHECKING:
+    from genelab.contracts import EnvContext
+    from genelab.managers.scene_entity_cfg import SceneEntityCfg
+
+
+def kinetic_energy_l2(env: "EnvContext", asset_cfg: "SceneEntityCfg | None" = None) -> torch.Tensor:
+    """``T²`` per env, where ``T`` is the articulation's kinetic energy.
+
+    Reads ``gs_handle.get_kinetic_energy()`` (Genesis 1.0) and squares the
+    per-env scalar. Non-negative; pair with a negative weight to discourage
+    high-energy motion.
+    """
+    ke = _asset_handle(env, asset_cfg).get_kinetic_energy()
+    return ke * ke
+
+
+def potential_energy(env: "EnvContext", asset_cfg: "SceneEntityCfg | None" = None) -> torch.Tensor:
+    """Signed potential energy of the named articulation, in Joules.
+
+    Reads ``gs_handle.get_potential_energy()`` directly. Sign follows
+    Genesis: bodies at world-frame z < 0 carry negative potential, bodies
+    above z = 0 carry positive potential. Pair with the appropriate signed
+    weight for the task.
+    """
+    return _asset_handle(env, asset_cfg).get_potential_energy()
+
+
+def energy_budget(
+    env: "EnvContext",
+    *,
+    target_total: float = 0.0,
+    asset_cfg: "SceneEntityCfg | None" = None,
+) -> torch.Tensor:
+    """``(total − target_total)²`` per env — squared deviation from a target total energy.
+
+    Reads ``gs_handle.get_total_energy()`` (= kinetic + potential) and squares
+    the deviation from ``target_total`` (in Joules). Non-negative; pair with
+    a negative weight to keep the articulation's total mechanical energy near
+    the target.
+    """
+    total = _asset_handle(env, asset_cfg).get_total_energy()
+    delta = total - target_total
+    return delta * delta

tests/test_mdp_rewards_energy.py

@@ -0,0 +1,74 @@
+"""Unit tests for ``genelab.mdp.rewards.energy``."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from genelab.mdp.rewards.energy import (  # noqa: E402
+    energy_budget,
+    kinetic_energy_l2,
+    potential_energy,
+)
+
+
+@dataclass
+class _FakeGsHandle:
+    """Minimal stand-in for a Genesis robot handle with the energy methods rewards read."""
+
+    kinetic: torch.Tensor
+    potential: torch.Tensor
+
+    def get_kinetic_energy(self) -> torch.Tensor:
+        return self.kinetic
+
+    def get_potential_energy(self) -> torch.Tensor:
+        return self.potential
+
+    def get_total_energy(self) -> torch.Tensor:
+        return self.kinetic + self.potential
+
+
+class _FakeArticulation:
+    def __init__(self, handle: _FakeGsHandle) -> None:
+        self.gs_handle = handle
+
+
+class _FakeEnv:
+    def __init__(self, articulations: dict[str, _FakeArticulation]) -> None:
+        self.articulations = articulations
+        # ``asset_handle`` falls back to ``env.robot`` when ``asset_cfg`` is None and
+        # the entity isn't in ``articulations``; expose the primary by default.
+        self.robot: Any = articulations["robot"].gs_handle
+
+
+def _env(*, kinetic: torch.Tensor, potential: torch.Tensor) -> _FakeEnv:
+    return _FakeEnv({"robot": _FakeArticulation(_FakeGsHandle(kinetic, potential))})
+
+
+def test_kinetic_energy_l2_squares_per_env() -> None:
+    env = _env(kinetic=torch.tensor([0.0, 1.5, -2.0]), potential=torch.zeros(3))
+    reward = kinetic_energy_l2(env)  # type: ignore[arg-type]
+    assert torch.allclose(reward, torch.tensor([0.0, 2.25, 4.0]))
+
+
+def test_potential_energy_returns_signed_value() -> None:
+    env = _env(kinetic=torch.zeros(2), potential=torch.tensor([-9.81, 4.2]))
+    reward = potential_energy(env)  # type: ignore[arg-type]
+    assert torch.allclose(reward, torch.tensor([-9.81, 4.2]))
+
+
+def test_energy_budget_squares_deviation_from_target() -> None:
+    env = _env(kinetic=torch.tensor([1.0, 2.0]), potential=torch.tensor([2.0, -1.0]))
+    # total = (3.0, 1.0); target=2.0 → deltas (1.0, -1.0) → squared (1.0, 1.0)
+    reward = energy_budget(env, target_total=2.0)  # type: ignore[arg-type]
+    assert torch.allclose(reward, torch.tensor([1.0, 1.0]))
+
+
+def test_energy_budget_default_target_zero() -> None:
+    env = _env(kinetic=torch.tensor([1.0, 3.0]), potential=torch.tensor([0.0, -1.0]))
+    # total = (1.0, 2.0); target=0.0 → squared (1.0, 4.0)
+    reward = energy_budget(env)  # type: ignore[arg-type]
+    assert torch.allclose(reward, torch.tensor([1.0, 4.0]))