diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d9a70e1..f838232d 100644 --- a/CHANGELOG.md +++ b/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 `` 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 diff --git a/docs/concepts/actuators.zh.md b/docs/concepts/actuators.zh.md index 4e84a708..c49dea33 100644 --- a/docs/concepts/actuators.zh.md +++ b/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。 diff --git a/src/genelab/actuator/__init__.py b/src/genelab/actuator/__init__.py index 15c56060..22c1c71c 100644 --- a/src/genelab/actuator/__init__.py +++ b/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", ] diff --git a/src/genelab/actuator/mujoco_style.py b/src/genelab/actuator/mujoco_style.py new file mode 100644 index 00000000..15a6d69d --- /dev/null +++ b/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 diff --git a/src/genelab/mdp/rewards/__init__.py b/src/genelab/mdp/rewards/__init__.py index 1599800f..4e4109b3 100644 --- a/src/genelab/mdp/rewards/__init__.py +++ b/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,6 +68,7 @@ "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", @@ -69,6 +76,7 @@ "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", diff --git a/src/genelab/mdp/rewards/energy.py b/src/genelab/mdp/rewards/energy.py new file mode 100644 index 00000000..e8c6da32 --- /dev/null +++ b/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 diff --git a/tests/test_mdp_rewards_energy.py b/tests/test_mdp_rewards_energy.py new file mode 100644 index 00000000..d1f5f4d9 --- /dev/null +++ b/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])) diff --git a/tests/test_mujoco_style_actuator.py b/tests/test_mujoco_style_actuator.py new file mode 100644 index 00000000..8fc47484 --- /dev/null +++ b/tests/test_mujoco_style_actuator.py @@ -0,0 +1,78 @@ +"""Unit tests for :class:`genelab.actuator.MujocoStyleActuatorCfg`. + +The cfg's only job is translating Mujoco's ``(gear, gain_prm, bias_prm)`` triple +into the underlying ``IdealPDActuatorCfg`` knobs and rejecting unsupported +``dyntype`` / ``gaintype`` / ``biastype`` combinations up front. Tests cover the +translation arithmetic and each rejection path. +""" + +import pytest + +torch = pytest.importorskip("torch") + +from genelab.actuator import IdealPDActuator, MujocoStyleActuatorCfg # noqa: E402 + + +def _cfg(**overrides: object) -> MujocoStyleActuatorCfg: + base: dict[str, object] = { + "target_names_expr": (".*",), + "gear": 1.0, + "bias_prm": (0.0, -100.0, -2.0), + "effort_limit": 80.0, + } + base.update(overrides) + return MujocoStyleActuatorCfg(**base) # type: ignore[arg-type] + + +def test_translation_to_stiffness_damping() -> None: + """``stiffness = -gear * bias_prm[1]``, ``damping = -gear * bias_prm[2]``.""" + cfg = _cfg(gear=1.5, bias_prm=(0.0, -200.0, -3.0)) + # -1.5 * -200 = 300; -1.5 * -3 = 4.5 + assert cfg.stiffness == pytest.approx(300.0) + assert cfg.damping == pytest.approx(4.5) + assert cfg.effort_limit == pytest.approx(80.0) + assert cfg.class_type is IdealPDActuator + + +def test_gear_one_is_identity_for_bias_prm() -> None: + cfg = _cfg(gear=1.0, bias_prm=(0.0, -50.0, -1.0)) + assert cfg.stiffness == pytest.approx(50.0) + assert cfg.damping == pytest.approx(1.0) + + +def test_non_zero_constant_bias_raises() -> None: + """``bias_prm[0]`` (constant force offset) has no IdealPD analogue.""" + with pytest.raises(ValueError, match="constant bias"): + _cfg(bias_prm=(0.5, -100.0, -2.0)) + + +def test_unsupported_dyntype_raises() -> None: + with pytest.raises(ValueError, match="dyntype"): + _cfg(dyntype="integrator") # type: ignore[arg-type] + + +def test_unsupported_gaintype_raises() -> None: + with pytest.raises(ValueError, match="gaintype"): + _cfg(gaintype="muscle") # type: ignore[arg-type] + + +def test_unsupported_biastype_raises() -> None: + with pytest.raises(ValueError, match="biastype"): + _cfg(biastype="muscle") # type: ignore[arg-type] + + +def test_explicit_stiffness_or_damping_raises() -> None: + """The cfg computes these from ``gear / bias_prm`` — direct overrides are confusing.""" + with pytest.raises(ValueError, match="stiffness / damping"): + _cfg(stiffness=10.0) # type: ignore[arg-type] + with pytest.raises(ValueError, match="stiffness / damping"): + _cfg(damping=0.5) # type: ignore[arg-type] + + +def test_preserves_passthrough_fields() -> None: + """``effort_limit`` / ``velocity_limit`` / ``armature`` / ``friction`` pass through unchanged.""" + cfg = _cfg(effort_limit=50.0, velocity_limit=12.0, armature=0.05, friction=0.1) + assert cfg.effort_limit == pytest.approx(50.0) + assert cfg.velocity_limit == pytest.approx(12.0) + assert cfg.armature == pytest.approx(0.05) + assert cfg.friction == pytest.approx(0.1)