Expose qfrc_actuator and qfrc_external on EntityData

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)

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
 -------------------------------
 

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.
 

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.

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 = """
+<mujoco>
+  <worldbody>
+    <body name="base">
+      <body name="upper" pos="0 0 0.5">
+        <joint name="shoulder" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" fromto="0 0 0 0 0 0.4" size="0.03" mass="1.0"/>
+        <body name="lower" pos="0 0 0.4">
+          <joint name="elbow" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" fromto="0 0 0 0 0 0.3" size="0.02" mass="0.5"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+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
+  )

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)