From 65ba706bc84eb99a2e9df7be190696de4a673c1d Mon Sep 17 00:00:00 2001 From: etaoxing Date: Thu, 2 Oct 2025 13:30:15 -0400 Subject: [PATCH] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` --- newton/_src/sim/builder.py | 6 ++++ newton/_src/sim/model.py | 4 +++ newton/_src/sim/state.py | 3 ++ newton/_src/solvers/mujoco/solver_mujoco.py | 40 +++++++++++++++++++-- 4 files changed, 51 insertions(+), 2 deletions(-) diff --git a/newton/_src/sim/builder.py b/newton/_src/sim/builder.py index 43e750d69..498197760 100644 --- a/newton/_src/sim/builder.py +++ b/newton/_src/sim/builder.py @@ -417,6 +417,7 @@ def __init__(self, up_axis: AxisType = Axis.Z, gravity: float = -9.81): self.joint_X_c = [] # frame of child com (in child coordinates) (constant) self.joint_q = [] self.joint_qd = [] + self.joint_act = [] self.joint_f = [] self.joint_type = [] @@ -1101,6 +1102,7 @@ def transform_mul(a, b): "joint_dof_mode", "joint_key", "joint_qd", + "joint_act", "joint_f", "joint_target", "joint_limit_lower", @@ -1336,6 +1338,7 @@ def add_axis_dim(dim: ModelBuilder.JointDofConfig): self.joint_q.append(0.0) for _ in range(dof_count): self.joint_qd.append(0.0) + self.joint_act.append(0.0) self.joint_f.append(0.0) if joint_type == JointType.FREE or joint_type == JointType.DISTANCE or joint_type == JointType.BALL: @@ -2084,6 +2087,7 @@ def collapse_fixed_joints(self, verbose=wp.config.verbose): "type": self.joint_type[i], "q": self.joint_q[q_start : q_start + q_dim], "qd": self.joint_qd[qd_start : qd_start + qd_dim], + "act": self.joint_act[qd_start : qd_start + qd_dim], "armature": self.joint_armature[qd_start : qd_start + qd_dim], "q_start": q_start, "qd_start": qd_start, @@ -2277,6 +2281,7 @@ def dfs(parent_body: int, child_body: int, incoming_xform: wp.transform, last_dy self.joint_child.clear() self.joint_q.clear() self.joint_qd.clear() + self.joint_act.clear() self.joint_q_start.clear() self.joint_qd_start.clear() self.joint_enabled.clear() @@ -4331,6 +4336,7 @@ def finalize(self, device: Devicelike | None = None, requires_grad: bool = False m.joint_target_kd = wp.array(self.joint_target_kd, dtype=wp.float32, requires_grad=requires_grad) m.joint_dof_mode = wp.array(self.joint_dof_mode, dtype=wp.int32) m.joint_target = wp.array(self.joint_target, dtype=wp.float32, requires_grad=requires_grad) + m.joint_act = wp.array(self.joint_act, dtype=wp.float32, requires_grad=requires_grad) m.joint_f = wp.array(self.joint_f, dtype=wp.float32, requires_grad=requires_grad) m.joint_effort_limit = wp.array(self.joint_effort_limit, dtype=wp.float32, requires_grad=requires_grad) m.joint_velocity_limit = wp.array(self.joint_velocity_limit, dtype=wp.float32, requires_grad=requires_grad) diff --git a/newton/_src/sim/model.py b/newton/_src/sim/model.py index f689e402b..db5efd556 100644 --- a/newton/_src/sim/model.py +++ b/newton/_src/sim/model.py @@ -226,6 +226,8 @@ def __init__(self, device: Devicelike | None = None): """Generalized joint positions for state initialization, shape [joint_coord_count], float.""" self.joint_qd = None """Generalized joint velocities for state initialization, shape [joint_dof_count], float.""" + self.joint_act = None + """Generalized joint actuation for state initialization, shape [joint_dof_count], float.""" self.joint_f = None """Generalized joint forces for state initialization, shape [joint_dof_count], float.""" self.joint_target = None @@ -405,6 +407,7 @@ def __init__(self, device: Devicelike | None = None): # attributes per joint dof self.attribute_frequency["joint_qd"] = "joint_dof" + self.attribute_frequency["joint_act"] = "joint_dof" self.attribute_frequency["joint_f"] = "joint_dof" self.attribute_frequency["joint_armature"] = "joint_dof" self.attribute_frequency["joint_target"] = "joint_dof" @@ -470,6 +473,7 @@ def state(self, requires_grad: bool | None = None) -> State: if self.joint_count: s.joint_q = wp.clone(self.joint_q, requires_grad=requires_grad) s.joint_qd = wp.clone(self.joint_qd, requires_grad=requires_grad) + s.joint_act = wp.clone(self.joint_act, requires_grad=requires_grad) return s diff --git a/newton/_src/sim/state.py b/newton/_src/sim/state.py index 2ff1e2171..6eafd6ec5 100644 --- a/newton/_src/sim/state.py +++ b/newton/_src/sim/state.py @@ -61,6 +61,9 @@ def __init__(self) -> None: self.joint_qd: wp.array | None = None """Generalized joint velocity coordinates, shape (joint_dof_count,), dtype float.""" + self.joint_act: wp.array | None = None + """Generalized joint actuation, shape (joint_dof_count,), dtype float.""" + def clear_forces(self) -> None: """ Clear all force arrays (for particles and bodies) in the state object. diff --git a/newton/_src/solvers/mujoco/solver_mujoco.py b/newton/_src/solvers/mujoco/solver_mujoco.py index 4a62e6be6..5209fe808 100644 --- a/newton/_src/solvers/mujoco/solver_mujoco.py +++ b/newton/_src/solvers/mujoco/solver_mujoco.py @@ -336,6 +336,7 @@ def convert_newton_contacts_to_mjwarp_kernel( def convert_mj_coords_to_warp_kernel( qpos: wp.array2d(dtype=wp.float32), qvel: wp.array2d(dtype=wp.float32), + qact: wp.array2d(dtype=wp.float32), joints_per_env: int, up_axis: int, joint_type: wp.array(dtype=wp.int32), @@ -345,6 +346,7 @@ def convert_mj_coords_to_warp_kernel( # outputs joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), + joint_act: wp.array(dtype=wp.float32), ): worldid, jntid = wp.tid() @@ -383,6 +385,14 @@ def convert_mj_coords_to_warp_kernel( joint_qd[wqd_i + 3] = w[0] joint_qd[wqd_i + 4] = w[1] joint_qd[wqd_i + 5] = w[2] + + joint_act[wqd_i + 0] = qact[worldid, q_i + 0] + joint_act[wqd_i + 1] = qact[worldid, q_i + 1] + joint_act[wqd_i + 2] = qact[worldid, q_i + 2] + joint_act[wqd_i + 3] = qact[worldid, q_i + 3] + joint_act[wqd_i + 4] = qact[worldid, q_i + 4] + joint_act[wqd_i + 5] = qact[worldid, q_i + 5] + elif type == JointType.BALL: # change quaternion order from wxyz to xyzw rot = wp.quat( @@ -398,6 +408,8 @@ def convert_mj_coords_to_warp_kernel( for i in range(3): # convert velocity components joint_qd[wqd_i + i] = qvel[worldid, qd_i + i] + # convert act components + joint_act[wqd_i + i] = qact[worldid, qd_i + i] else: axis_count = joint_dof_dim[jntid, 0] + joint_dof_dim[jntid, 1] for i in range(axis_count): @@ -406,12 +418,15 @@ def convert_mj_coords_to_warp_kernel( for i in range(axis_count): # convert velocity components joint_qd[wqd_i + i] = qvel[worldid, qd_i + i] + # convert act components + joint_act[wqd_i + i] = qact[worldid, qd_i + i] @wp.kernel def convert_warp_coords_to_mj_kernel( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), + joint_act: wp.array(dtype=wp.float32), joints_per_env: int, up_axis: int, joint_type: wp.array(dtype=wp.int32), @@ -421,6 +436,7 @@ def convert_warp_coords_to_mj_kernel( # outputs qpos: wp.array2d(dtype=wp.float32), qvel: wp.array2d(dtype=wp.float32), + qact: wp.array2d(dtype=wp.float32), ): worldid, jntid = wp.tid() @@ -460,6 +476,13 @@ def convert_warp_coords_to_mj_kernel( qvel[worldid, qd_i + 4] = w[1] qvel[worldid, qd_i + 5] = w[2] + qact[worldid, qd_i + 0] = joint_act[wqd_i + 0] + qact[worldid, qd_i + 1] = joint_act[wqd_i + 1] + qact[worldid, qd_i + 2] = joint_act[wqd_i + 2] + qact[worldid, qd_i + 3] = joint_act[wqd_i + 3] + qact[worldid, qd_i + 4] = joint_act[wqd_i + 4] + qact[worldid, qd_i + 5] = joint_act[wqd_i + 5] + elif type == JointType.BALL: # change quaternion order from xyzw to wxyz qpos[worldid, q_i + 0] = joint_q[wq_i + 1] @@ -469,6 +492,8 @@ def convert_warp_coords_to_mj_kernel( for i in range(3): # convert velocity components qvel[worldid, qd_i + i] = joint_qd[wqd_i + i] + # convert act components + qact[worldid, qd_i + i] = joint_act[wqd_i + i] else: axis_count = joint_dof_dim[jntid, 0] + joint_dof_dim[jntid, 1] for i in range(axis_count): @@ -477,6 +502,8 @@ def convert_warp_coords_to_mj_kernel( for i in range(axis_count): # convert velocity components qvel[worldid, qd_i + i] = joint_qd[wqd_i + i] + # convert act components + qact[worldid, qd_i + i] = joint_act[wqd_i + i] @wp.kernel @@ -1499,18 +1526,22 @@ def update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: Sta # we have an MjWarp Data object qpos = mj_data.qpos qvel = mj_data.qvel + qact = mj_data.qact nworld = mj_data.nworld else: # we have an MjData object from Mujoco qpos = wp.empty((1, model.joint_coord_count), dtype=wp.float32, device=model.device) qvel = wp.empty((1, model.joint_dof_count), dtype=wp.float32, device=model.device) + qact = wp.empty((1, model.joint_dof_count), dtype=wp.float32, device=model.device) nworld = 1 if state is None: joint_q = model.joint_q joint_qd = model.joint_qd + joint_act = model.joint_act else: joint_q = state.joint_q joint_qd = state.joint_qd + joint_act = state.joint_act joints_per_env = model.joint_count // nworld wp.launch( convert_warp_coords_to_mj_kernel, @@ -1518,6 +1549,7 @@ def update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: Sta inputs=[ joint_q, joint_qd, + joint_act, joints_per_env, model.up_axis, model.joint_type, @@ -1525,12 +1557,13 @@ def update_mjc_data(self, mj_data: MjWarpData | MjData, model: Model, state: Sta model.joint_qd_start, model.joint_dof_dim, ], - outputs=[qpos, qvel], + outputs=[qpos, qvel, qact], device=model.device, ) if not is_mjwarp: mj_data.qpos[:] = qpos.numpy().flatten()[: len(mj_data.qpos)] mj_data.qvel[:] = qvel.numpy().flatten()[: len(mj_data.qvel)] + mj_data.qfrc_actuator[:] = qact.numpy().flatten()[: len(mj_data.qfrc_actuator)] def update_newton_state( self, @@ -1544,6 +1577,7 @@ def update_newton_state( # we have an MjWarp Data object qpos = mj_data.qpos qvel = mj_data.qvel + qact = mj_data.qfrc_actuator nworld = mj_data.nworld xpos = mj_data.xpos @@ -1552,6 +1586,7 @@ def update_newton_state( # we have an MjData object from Mujoco qpos = wp.array([mj_data.qpos], dtype=wp.float32, device=model.device) qvel = wp.array([mj_data.qvel], dtype=wp.float32, device=model.device) + qact = wp.array([mj_data.qfrc_actuator], dtype=wp.float32, device=model.device) nworld = 1 xpos = wp.array([mj_data.xpos], dtype=wp.vec3, device=model.device) @@ -1563,6 +1598,7 @@ def update_newton_state( inputs=[ qpos, qvel, + qact, joints_per_env, int(model.up_axis), model.joint_type, @@ -1570,7 +1606,7 @@ def update_newton_state( model.joint_qd_start, model.joint_dof_dim, ], - outputs=[state.joint_q, state.joint_qd], + outputs=[state.joint_q, state.joint_qd, state.joint_act], device=model.device, )