diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index bccf1843d..c5154723b 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -17,6 +17,11 @@ Added the dump viewer for fixed-base entities. - Implemented ``ActionTermCfg.clip`` for clamping processed actions after scale and offset (:issue:`771`). +- Added ``qfrc_actuator`` and ``qfrc_external`` generalized force accessors + to ``EntityData``. ``qfrc_actuator`` gives actuator forces in joint space + (projected through the transmission). ``qfrc_external`` recovers the + generalized force from body external wrenches (``xfrc_applied``) + (:issue:`776`). Changed ^^^^^^^ @@ -27,10 +32,17 @@ Changed ``clear_state`` only zeroed actuator targets without resetting actuator internal state (e.g. delay buffers), which could cause stale commands after teleporting the robot to a new pose. +- Removed ``EntityData.generalized_force``. The property was bugged (indexed + free joint DOFs instead of articulated DOFs) and the name was ambiguous. + Use ``qfrc_actuator`` or ``qfrc_external`` instead (:issue:`776`). Fixed ^^^^^ +- ``electrical_power_cost`` now uses ``qfrc_actuator`` (joint space) instead + of ``actuator_force`` (actuation space) for mechanical power computation. + Previously the reward was incorrect for actuators with gear ratios other + than 1 (:issue:`776`). - ``dr.pseudo_inertia`` no longer loads cuSOLVER, eliminating ~4 GB of persistent GPU memory overhead. Cholesky and eigendecomposition are now computed analytically for the small matrices involved (4x4 and 3x3) diff --git a/docs/source/entity/entity_data.rst b/docs/source/entity/entity_data.rst index 687c203ba..e554dbab8 100644 --- a/docs/source/entity/entity_data.rst +++ b/docs/source/entity/entity_data.rst @@ -258,13 +258,39 @@ Joint properties cover 1-DOF revolute and prismatic joints. The free joint - Joint accelerations in rad/s² or m/s² * - ``actuator_force`` - ``[num_envs, num_actuators]`` - - Scalar actuation force in actuation space. Use this in reward and - observation terms where joint torques are needed. - * - ``generalized_force`` - - ``[num_envs, nv]`` - - Generalized forces applied to the free-joint DOFs + - Scalar actuator output in actuation space (per actuator). This is + the force before projection through the transmission Jacobian. For + actuator forces in joint space, use ``qfrc_actuator`` instead. +.. _generalized-forces: + +Reference: generalized forces +----------------------------- + +These properties expose selected components of MuJoCo's generalized +force decomposition, sliced to this entity's articulated joint DOFs. +Free joint DOFs are excluded. All shapes are ``[num_envs, nv]`` where +``nv`` is the number of articulated DOFs belonging to this entity. + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Property + - Description + * - ``qfrc_actuator`` + - Forces produced by all actuators, mapped into joint space. For + motors this is the commanded torque times the gear ratio. For + position and velocity actuators this is the force computed by + the internal PD law. When ``actuatorgravcomp`` is enabled on a + joint, the gravity compensation force is included here. + * - ``qfrc_external`` + - Forces on joints due to Cartesian wrenches applied to bodies + via ``xfrc_applied``. This is the :math:`J^\top F` mapping. + MuJoCo does not store this term separately; the property + recovers it from other force components after ``forward()``. + Reference: geom and site state ------------------------------- diff --git a/src/mjlab/entity/data.py b/src/mjlab/entity/data.py index 80bef750b..54f18d48e 100644 --- a/src/mjlab/entity/data.py +++ b/src/mjlab/entity/data.py @@ -234,6 +234,11 @@ def _resolve_env_ids( return env_ids[:, None] return env_ids + def _joint_dof_field(self, field_name: str) -> torch.Tensor: + """Return a generalized-force field sliced to this entity's joint DoFs.""" + field = getattr(self.data, field_name) + return field[:, self.indexing.joint_v_adr] + # Root properties @property @@ -387,24 +392,55 @@ def tendon_vel(self) -> torch.Tensor: """Tendon velocities. Shape (num_envs, num_tendons).""" return self.data.ten_velocity[:, self.indexing.tendon_ids] + # Generalized forces + @property def joint_torques(self) -> torch.Tensor: """Joint torques. Shape (num_envs, nv).""" raise NotImplementedError( - "Joint torques are not currently available. " - "Consider using 'actuator_force' property for actuation forces, " - "or 'generalized_force' property for generalized forces applied to the DoFs." + "Joint torques are ambiguous. Use 'qfrc_actuator' for actuator forces " + "in joint space, or 'qfrc_external' for body wrench contributions." ) @property def actuator_force(self) -> torch.Tensor: - """Scalar actuation force in actuation space. Shape (num_envs, nu).""" + """Scalar actuator output in actuation space. Shape (num_envs, nu). + + This is not the same as joint-space generalized force. Use ``qfrc_actuator`` for + the actuator contribution projected into DoF space. + """ return self.data.actuator_force[:, self.indexing.ctrl_ids] @property - def generalized_force(self) -> torch.Tensor: - """Generalized forces applied to the DoFs. Shape (num_envs, nv).""" - return self.data.qfrc_applied[:, self.indexing.free_joint_v_adr] + def qfrc_actuator(self) -> torch.Tensor: + """Forces produced by all actuators, mapped into joint space. + + For motors this is the commanded torque times the gear ratio. For position and + velocity actuators this is the force computed by the internal PD law. When + ``actuatorgravcomp`` is enabled on a joint, the gravity compensation force is + included here. + """ + return self._joint_dof_field("qfrc_actuator") + + @property + def qfrc_external(self) -> torch.Tensor: + """Forces on joints due to Cartesian wrenches applied to bodies. + + When a force or torque is applied to a body via ``xfrc_applied``, this property + gives the equivalent joint forces (the Jacobian transpose mapping). + """ + # MuJoCo folds J^T * xfrc_applied into qfrc_smooth without storing it. + # Recover via the qfrc_smooth identity: + # qfrc_smooth = qfrc_actuator + qfrc_passive - qfrc_bias + # + qfrc_applied + J^T * xfrc_applied + f = self._joint_dof_field + return ( + f("qfrc_smooth") + - f("qfrc_actuator") + - f("qfrc_applied") + - f("qfrc_passive") + + f("qfrc_bias") + ) # Pose and velocity component accessors. diff --git a/src/mjlab/envs/mdp/rewards.py b/src/mjlab/envs/mdp/rewards.py index 966333435..8712d837b 100644 --- a/src/mjlab/envs/mdp/rewards.py +++ b/src/mjlab/envs/mdp/rewards.py @@ -139,15 +139,11 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv): joint_ids, _ = asset.find_joints( cfg.params["asset_cfg"].joint_names, ) - actuator_ids, _ = asset.find_actuators( - cfg.params["asset_cfg"].joint_names, - ) self._joint_ids = torch.tensor(joint_ids, device=env.device, dtype=torch.long) - self._actuator_ids = torch.tensor(actuator_ids, device=env.device, dtype=torch.long) def __call__(self, env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: asset: Entity = env.scene[asset_cfg.name] - tau = asset.data.actuator_force[:, self._actuator_ids] + tau = asset.data.qfrc_actuator[:, self._joint_ids] qd = asset.data.joint_vel[:, self._joint_ids] mech = tau * qd mech_pos = torch.clamp(mech, min=0.0) # Don't penalize regen. diff --git a/tests/test_entity_data.py b/tests/test_entity_data.py index 3b581ca99..0da1a2317 100644 --- a/tests/test_entity_data.py +++ b/tests/test_entity_data.py @@ -1,10 +1,17 @@ """Tests for EntityData.""" import mujoco +import numpy as np import pytest import torch -from conftest import get_test_device +from conftest import ( + create_entity_with_actuator, + get_test_device, + initialize_entity, + load_fixture_xml, +) +from mjlab.actuator import BuiltinMotorActuatorCfg from mjlab.entity import Entity, EntityCfg from mjlab.sim.sim import Simulation, SimulationCfg @@ -282,3 +289,109 @@ def test_entity_data_reset_partial_envs(device): assert torch.all(entity.data.tendon_effort_target[1] == 0.0) assert torch.all(entity.data.tendon_effort_target[2] == 9.0) assert torch.all(entity.data.tendon_effort_target[3] == 0.0) + + +# Generalized force accessor tests. + +TWO_LINK_ARM_XML = """ + + + + + + + + + + + + + + +""" + + +def test_qfrc_actuator_slices_joint_dofs_only(device): + """qfrc_actuator should expose only articulated joint DoFs.""" + xml = load_fixture_xml("floating_base_articulated") + cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml)) + entity = Entity(cfg) + entity, sim = initialize_entity_with_sim(entity, device) + + joint_v_adr = entity.indexing.joint_v_adr + free_v_adr = entity.indexing.free_joint_v_adr + nv = sim.data.qvel.shape[1] + + values = torch.arange(1, nv + 1, device=device, dtype=torch.float32) + sim.data.qfrc_actuator[0] = values + expected = values[joint_v_adr] + actual = entity.data.qfrc_actuator[0] + assert torch.equal(actual, expected) + + assert free_v_adr.numel() > 0 + assert entity.data.qfrc_actuator.shape[-1] == joint_v_adr.numel() + assert entity.data.qfrc_actuator.shape[-1] != free_v_adr.numel() + + +def test_qfrc_actuator_matches_motor_command_with_gear(device): + """qfrc_actuator should equal actuator_force projected through joint gear.""" + xml = load_fixture_xml("fixed_base_articulated") + entity = create_entity_with_actuator( + xml, + BuiltinMotorActuatorCfg( + target_names_expr=("joint.*",), + effort_limit=100.0, + gear=3.0, + ), + ) + entity, sim = initialize_entity(entity, device) + + commanded = torch.tensor([[2.0, -1.0]], device=device) + entity.set_joint_effort_target(commanded) + entity.write_data_to_sim() + sim.forward() + + assert torch.allclose(entity.data.actuator_force, commanded, atol=1e-6) + assert torch.allclose(entity.data.qfrc_actuator, 3.0 * commanded, atol=1e-6) + + +def test_qfrc_external_matches_mj_applyFT(device): + """qfrc_external should match J^T * xfrc_applied computed by CPU MuJoCo.""" + cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(TWO_LINK_ARM_XML)) + entity = Entity(cfg) + entity, sim = initialize_entity_with_sim(entity, device) + + entity.write_joint_state_to_sim( + torch.tensor([[0.5, -0.3]], device=device), + torch.tensor([[0.0, 0.0]], device=device), + ) + sim.forward() + + lower_body_id = int(entity.indexing.body_ids[1].item()) + force = torch.tensor([1.0, 0.0, -2.0], device=device) + torque = torch.tensor([0.0, 0.5, 0.0], device=device) + sim.data.xfrc_applied[0, lower_body_id, 0:3] = force + sim.data.xfrc_applied[0, lower_body_id, 3:6] = torque + sim.forward() + + mjd = mujoco.MjData(sim.mj_model) + mjd.qpos[:] = sim.data.qpos[0].cpu().numpy() + mjd.qvel[:] = sim.data.qvel[0].cpu().numpy() + mujoco.mj_forward(sim.mj_model, mjd) + + qfrc_reference = np.zeros(sim.mj_model.nv) + mujoco.mj_applyFT( + sim.mj_model, + mjd, + force.cpu().numpy().astype(np.float64), + torque.cpu().numpy().astype(np.float64), + mjd.xipos[lower_body_id], + lower_body_id, + qfrc_reference, + ) + expected = torch.from_numpy(qfrc_reference).to(device=device, dtype=torch.float32) + + joint_v_adr = entity.indexing.joint_v_adr + assert torch.allclose( + entity.data.qfrc_external, expected[joint_v_adr].unsqueeze(0), atol=1e-5 + ) diff --git a/tests/test_rewards.py b/tests/test_rewards.py index bd3708bc7..adeb4119a 100644 --- a/tests/test_rewards.py +++ b/tests/test_rewards.py @@ -200,14 +200,13 @@ def test_electrical_power_cost_partially_actuated(device): reward = electrical_power_cost(reward_cfg, env) assert len(reward._joint_ids) == 2 - assert len(reward._actuator_ids) == 2 # Test case 1: All forces and velocities aligned (all positive work). - # actuated_joint1 (qvel[:, 0]), actuated_joint2 (qvel[:, 2]). - # Note: qvel[:, 1] is the passive joint, not used in power calculation. - sim.data.actuator_force[:] = torch.tensor( - [[2.0, 3.0], [1.0, 4.0]], device=device, dtype=torch.float32 - ) + # actuated_joint1 (dof 0), actuated_joint2 (dof 2). + # Note: dof 1 is the passive joint, not used in power calculation. + sim.data.qfrc_actuator[:] = 0.0 + sim.data.qfrc_actuator[:, 0] = torch.tensor([2.0, 1.0], device=device) + sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, 4.0], device=device) sim.data.qvel[:] = 0.0 sim.data.qvel[:, 0] = torch.tensor([1.0, 2.0], device=device) sim.data.qvel[:, 2] = torch.tensor([2.0, 1.0], device=device) @@ -219,9 +218,9 @@ def test_electrical_power_cost_partially_actuated(device): assert torch.allclose(power_cost, expected) # Test case 2: Some negative forces (regenerative braking, not penalized). - sim.data.actuator_force[:] = torch.tensor( - [[-2.0, 3.0], [1.0, -4.0]], device=device, dtype=torch.float32 - ) + sim.data.qfrc_actuator[:] = 0.0 + sim.data.qfrc_actuator[:, 0] = torch.tensor([-2.0, 1.0], device=device) + sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, -4.0], device=device) power_cost = reward(env, asset_cfg)