diff --git a/motrix_env_motrixsim/pyproject.toml b/motrix_env_motrixsim/pyproject.toml index 427106c0..342853bc 100644 --- a/motrix_env_motrixsim/pyproject.toml +++ b/motrix_env_motrixsim/pyproject.toml @@ -12,7 +12,7 @@ readme = "README.md" license = "Apache-2.0" dependencies = [ "motrix-env-core", - "motrixsim==0.10.1.dev123478", + "motrixsim==0.10.1.dev123565", "numpy>=1.26", ] diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py index 603598e9..dab8e14c 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py @@ -3,6 +3,7 @@ """MotrixSim compiler and executable program for declarative sim writes.""" +import math from collections.abc import Callable from dataclasses import dataclass from typing import Protocol @@ -45,10 +46,17 @@ def apply(self, dof_pos, dof_vel, buffers, env_ids) -> None: ... @dataclass(frozen=True) class _CompiledWrite: - op: _WriteOp + """One declared write: either native plan fields or a numpy scatter op.""" + + op: _WriteOp | None = None reset_op: _ResetPatchOp | None = None + fields: tuple[tuple[str, mtx.write.WriteSource], ...] = () + # Logical buffer shape synthesized across a multi-field declaration; the + # fields interleave per leading axis so one contiguous slice views them. + synthesized_shape: tuple[int, ...] | None = None pos_indices: np.ndarray | None = None vel_indices: np.ndarray | None = None + ctrl_indices: np.ndarray | None = None refresh_kinematics: bool = False @@ -70,63 +78,95 @@ def _begin_compile(self) -> None: self._pending = [] def _build_program(self, *, reset: bool, forward_kinematics: bool) -> WriteProgram: - buffers: dict[str, np.ndarray] = {} + fields: dict[str, mtx.write.WriteSource] = {} ops: list[tuple[_CompiledWrite, np.ndarray]] = [] + buffers: dict[str, np.ndarray] = {} ctrl_owners: dict[int, str] = {} claimed_pos: dict[int, str] = {} claimed_vel: dict[int, str] = {} - refresh_kinematics = False + lead_op_count = 0 + any_kinematics = False + numpy_kinematics = False for name, compiled in self._pending: - if isinstance(compiled.op, _CtrlOp): - self._claim_ctrl_targets(name, compiled.op.indices, ctrl_owners) + if compiled.ctrl_indices is not None: + self._claim_ctrl_targets(name, compiled.ctrl_indices, ctrl_owners) if compiled.pos_indices is not None: self._claim(name, compiled.pos_indices, claimed_pos, "position") if compiled.vel_indices is not None: self._claim(name, compiled.vel_indices, claimed_vel, "velocity") + any_kinematics |= compiled.refresh_kinematics + if compiled.fields: + if not fields: + lead_op_count = len(ops) + for field_name, source in compiled.fields: + fields[field_name] = source + continue + numpy_kinematics |= compiled.refresh_kinematics sub = compiled.op.alloc(self._data.shape[0]) ops.append((compiled, sub)) buffers[name] = sub - refresh_kinematics |= compiled.refresh_kinematics + native = None + manual_fk = False + needs_kinematics = forward_kinematics and (reset or any_kinematics) + synthesized: dict[str, tuple[int, tuple[int, ...]]] = {} + if fields: + # Forward kinematics runs inside the native execute unless numpy + # scatter ops must still land before the refresh; either way each + # execute refreshes at most once. + native = self._model.compile_write( + fields, reset=reset, forward_kinematic=needs_kinematics and not numpy_kinematics + ).allocate(self._data) + manual_fk = needs_kinematics and numpy_kinematics + offsets = {field.name: field.offset for field in native.fields} + for name, compiled in self._pending: + if compiled.synthesized_shape is not None: + first_field = compiled.fields[0][0] + synthesized[name] = (offsets[first_field], compiled.synthesized_shape) + else: + manual_fk = needs_kinematics return _MotrixSimWriteProgram( self._model, self._data, self._masked_rows, buffers, ops, + native=native, reset=reset, - refresh_kinematics=forward_kinematics and (reset or refresh_kinematics), + refresh_kinematics=manual_fk, + lead_op_count=lead_op_count, + synthesized=synthesized, ) def compile_dof_position(self, name: str, write: DofPositionWrite) -> None: del write indices = np.arange(self._model.num_dof_pos, dtype=np.int64) - op = _DofChannelOp(indices) + op = _DofChannelOp(self._model, indices) self._pending.append((name, _CompiledWrite(op, op, pos_indices=indices, refresh_kinematics=True))) def compile_dof_velocity(self, name: str, write: DofVelocityWrite) -> None: del write indices = np.arange(self._model.num_dof_vel, dtype=np.int64) - op = _DofChannelOp(indices, velocity=True) + op = _DofChannelOp(self._model, indices, velocity=True) self._pending.append((name, _CompiledWrite(op, op, vel_indices=indices))) def compile_body_joint_position(self, name: str, write: BodyJointPositionWrite) -> None: indices = np.asarray(_named_body(self._model, write.body).get_dof_pos_indices(False), dtype=np.int64) - op = _DofChannelOp(indices) + op = _DofChannelOp(self._model, indices) self._pending.append((name, _CompiledWrite(op, op, pos_indices=indices, refresh_kinematics=True))) def compile_body_joint_velocity(self, name: str, write: BodyJointVelocityWrite) -> None: indices = np.asarray(_named_body(self._model, write.body).get_dof_vel_indices(False), dtype=np.int64) - op = _DofChannelOp(indices, velocity=True) + op = _DofChannelOp(self._model, indices, velocity=True) self._pending.append((name, _CompiledWrite(op, op, vel_indices=indices))) def compile_joint_position(self, name: str, write: JointPositionWrite) -> None: indices = np.asarray([joint.dof_pos_index for joint in self._joints(name, write.joints)], dtype=np.int64) - op = _DofChannelOp(indices) + op = _DofChannelOp(self._model, indices) self._pending.append((name, _CompiledWrite(op, op, pos_indices=indices, refresh_kinematics=True))) def compile_joint_velocity(self, name: str, write: JointVelocityWrite) -> None: indices = np.asarray([joint.dof_vel_index for joint in self._joints(name, write.joints)], dtype=np.int64) - op = _DofChannelOp(indices, velocity=True) + op = _DofChannelOp(self._model, indices, velocity=True) self._pending.append((name, _CompiledWrite(op, op, vel_indices=indices))) def compile_ctrl_targets(self, name: str, write: CtrlTargetsWrite) -> None: @@ -140,74 +180,117 @@ def compile_ctrl_targets(self, name: str, write: CtrlTargetsWrite) -> None: indices = np.asarray( [_named_actuator(self._model, actuator).index for actuator in write.actuators], dtype=np.int64 ) - self._pending.append((name, _CompiledWrite(_CtrlOp(indices)))) + self._pending.append( + (name, _CompiledWrite(fields=((name, mtx.write.ActuatorCtrls(indices)),), ctrl_indices=indices)) + ) def compile_body_position(self, name: str, write: BodyPositionWrite) -> None: bases = self._floating_bases(name, write.bodies, type(write).__name__) indices = np.asarray([base.dof_pos_indices[:3] for base in bases], dtype=np.int64) - op = _MultiTargetOp(bases, "set_translation", 3) self._pending.append( ( name, - _CompiledWrite(op, _DofComponentPatchOp(indices), pos_indices=indices.ravel(), refresh_kinematics=True), + _CompiledWrite( + fields=((name, mtx.write.BodyPosition(list(write.bodies))),), + pos_indices=indices.ravel(), + refresh_kinematics=True, + ), ) ) def compile_body_rotation(self, name: str, write: BodyRotationWrite) -> None: bases = self._floating_bases(name, write.bodies, type(write).__name__) indices = np.asarray([base.dof_pos_indices[3:] for base in bases], dtype=np.int64) - op = _MultiTargetOp(bases, "set_rotation", 4, contiguous=True) self._pending.append( ( name, - _CompiledWrite(op, _DofComponentPatchOp(indices), pos_indices=indices.ravel(), refresh_kinematics=True), + _CompiledWrite( + fields=((name, mtx.write.BodyRotation(list(write.bodies))),), + pos_indices=indices.ravel(), + refresh_kinematics=True, + ), ) ) def compile_body_linear_velocity(self, name: str, write: BodyLinearVelocityWrite) -> None: bases = self._floating_bases(name, write.bodies, type(write).__name__) indices = np.asarray([base.dof_vel_indices[:3] for base in bases], dtype=np.int64) - op = _MultiTargetOp(bases, "set_global_linear_velocity", 3) self._pending.append( - (name, _CompiledWrite(op, _DofComponentPatchOp(indices, velocity=True), vel_indices=indices.ravel())) + ( + name, + _CompiledWrite( + fields=((name, mtx.write.BodyLinearVelocity(list(write.bodies))),), + vel_indices=indices.ravel(), + ), + ) ) def compile_body_angular_velocity(self, name: str, write: BodyAngularVelocityWrite) -> None: bases = self._floating_bases(name, write.bodies, type(write).__name__) indices = np.asarray([base.dof_vel_indices[3:] for base in bases], dtype=np.int64) - op = _MultiTargetOp(bases, "set_global_angular_velocity", 3) self._pending.append( - (name, _CompiledWrite(op, _DofComponentPatchOp(indices, velocity=True), vel_indices=indices.ravel())) + ( + name, + _CompiledWrite( + fields=((name, mtx.write.BodyAngularVelocity(list(write.bodies))),), + vel_indices=indices.ravel(), + ), + ) ) def compile_mocap_pose(self, name: str, write: MocapPoseWrite) -> None: bodies = self._targets(name, write.bodies, "body", _named_body) - mocaps = [] for body_name, body in zip(write.bodies, bodies): if body.mocap is None: raise ValueError(f"MocapPoseWrite body {body_name!r} is not a mocap body.") - mocaps.append(body.mocap) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(mocaps, "set_pose", 7), refresh_kinematics=True))) + # Per-body interleaved position/rotation fields keep the whole group + # contiguous in the program buffer, so buffer(name) synthesizes the + # neutral (N, B, 7) pose layout as one zero-copy slice. + fields = [] + for index, body_name in enumerate(write.bodies): + fields.append((f"{name}.pos.{index}", mtx.write.BodyPosition((body_name,)))) + fields.append((f"{name}.rot.{index}", mtx.write.BodyRotation((body_name,)))) + self._pending.append( + ( + name, + _CompiledWrite( + fields=tuple(fields), + synthesized_shape=(len(write.bodies), 7), + refresh_kinematics=True, + ), + ) + ) def compile_actuator_kp(self, name: str, write: ActuatorKpWrite) -> None: targets = self._targets(name, write.actuators, "actuator", _named_actuator) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(targets, "set_kp_override", 1)))) + indices = np.asarray([target.index for target in targets], dtype=np.int64) + self._pending.append((name, _CompiledWrite(fields=((name, mtx.write.ActuatorKpOverride(indices)),)))) def compile_actuator_damping(self, name: str, write: ActuatorDampingWrite) -> None: targets = self._targets(name, write.actuators, "actuator", _named_actuator) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(targets, "set_damping_override", 1)))) + indices = np.asarray([target.index for target in targets], dtype=np.int64) + self._pending.append((name, _CompiledWrite(fields=((name, mtx.write.ActuatorDampingOverride(indices)),)))) def compile_body_mass(self, name: str, write: BodyMassWrite) -> None: targets = self._targets(name, write.links, "link", _named_link) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(targets, "set_mass_override", 1)))) + self._pending.append( + (name, _CompiledWrite(fields=((name, mtx.write.LinkMassOverride([link.name for link in targets])),))) + ) def compile_body_com(self, name: str, write: BodyComWrite) -> None: targets = self._targets(name, write.links, "link", _named_link) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(targets, "set_center_of_mass_override", 3)))) + self._pending.append( + ( + name, + _CompiledWrite(fields=((name, mtx.write.LinkCenterOfMassOverride([link.name for link in targets])),)), + ) + ) def compile_geom_friction(self, name: str, write: GeomFrictionWrite) -> None: targets = self._targets(name, write.geoms, "geom", _named_geom) - self._pending.append((name, _CompiledWrite(_MultiTargetOp(targets, "set_friction_override", 3)))) + self._pending.append( + (name, _CompiledWrite(fields=((name, mtx.write.GeomFrictionOverride([geom.name for geom in targets])),))) + ) def _joints(self, name: str, joint_names: tuple[str, ...]): if not joint_names: @@ -260,7 +343,7 @@ def _claim_ctrl_targets(self, name: str, indices: np.ndarray, owners: dict[int, class _MotrixSimWriteProgram(WriteProgram): - """Compiled MotrixSim writes with program-owned value buffers.""" + """Compiled MotrixSim writes: one native FFI plan plus numpy scatter ops.""" def __init__( self, @@ -270,19 +353,31 @@ def __init__( buffers: dict[str, np.ndarray], ops: list[tuple[_CompiledWrite, np.ndarray]], *, + native: mtx.write.WriteProgram | None = None, reset: bool, refresh_kinematics: bool, + lead_op_count: int = 0, + synthesized: dict[str, tuple[int, tuple[int, ...]]] | None = None, ) -> None: self._model = model self._data = data self._masked_rows = masked_rows self._buffers = buffers self._ops = ops + self._native = native self._reset = reset self._refresh_kinematics = refresh_kinematics + self._lead_op_count = lead_op_count + self._synthesized = synthesized or {} def buffer(self, name: str) -> np.ndarray: - return self._buffers[name] + if name in self._buffers: + return self._buffers[name] + if name in self._synthesized: + offset, shape = self._synthesized[name] + view = self._native.buffer[:, offset : offset + math.prod(shape)] + return view.reshape(*self._data.shape, *shape) + return self._native[name] def execute(self, env_ids: np.ndarray | None = None) -> None: if env_ids is not None: @@ -297,38 +392,49 @@ def execute(self, env_ids: np.ndarray | None = None) -> None: selected_ids = np.arange(self._data.shape[0], dtype=np.int64) if env_ids is None else np.sort(env_ids) rows = self._data if env_ids is None else self._masked_rows(selected_ids) idx = slice(None) if env_ids is None else selected_ids + if self._native is None: + self._execute_numpy(rows, selected_ids, idx) + return + ids = None if env_ids is None else selected_ids + # Native fields apply before numpy ops declared after them; reset + # programs always run the native reset-and-fields pass first. + start = 0 if self._reset else self._lead_op_count + for compiled, sub_buffers in self._ops[:start]: + compiled.op(sub_buffers, idx, rows) + self._native.execute(self._data, ids) + for compiled, sub_buffers in self._ops[start:]: + compiled.op(sub_buffers, idx, rows) + if self._refresh_kinematics: + self._model.forward_kinematic(rows) + + def _execute_numpy(self, rows: mtx.SceneData, env_ids: np.ndarray, idx: np.ndarray | slice) -> None: if self._reset: - self._execute_reset(rows, selected_ids, idx) + self._execute_reset(rows, env_ids) return for compiled, sub_buffers in self._ops: compiled.op(sub_buffers, idx, rows) if self._refresh_kinematics: self._model.forward_kinematic(rows) - def _execute_reset(self, rows: mtx.SceneData, env_ids: np.ndarray, idx: np.ndarray | slice) -> None: + def _execute_reset(self, rows: mtx.SceneData, env_ids: np.ndarray) -> None: + # Every numpy op is a fused dof-channel patch: fold the declared + # values into the reset state and restore it in one native call. default_dof_pos = np.asarray(self._model.compute_init_dof_pos(), dtype=np.float32) dof_pos = np.broadcast_to(default_dof_pos, (env_ids.size, self._model.num_dof_pos)).copy() dof_vel = np.zeros((env_ids.size, self._model.num_dof_vel), dtype=np.float32) - post_reset_ops = [] for compiled, buffers in self._ops: - if compiled.reset_op is None: - post_reset_ops.append((compiled.op, buffers)) - else: - compiled.reset_op.apply(dof_pos, dof_vel, buffers, env_ids) - kwargs = {"forward_kinematic": self._refresh_kinematics and not post_reset_ops} + compiled.reset_op.apply(dof_pos, dof_vel, buffers, env_ids) + kwargs = {"forward_kinematic": self._refresh_kinematics} if self._model.num_dof_pos: kwargs["dof_pos"] = np.ascontiguousarray(dof_pos) if self._model.num_dof_vel: kwargs["dof_vel"] = np.ascontiguousarray(dof_vel) rows.reset(self._model, **kwargs) - for op, buffers in post_reset_ops: - op(buffers, idx, rows) - if post_reset_ops and self._refresh_kinematics: - self._model.forward_kinematic(rows) class _DofChannelOp: - def __init__(self, indices: np.ndarray, *, velocity: bool = False) -> None: + def __init__(self, model: mtx.SceneModel, indices: np.ndarray, *, velocity: bool = False) -> None: + self._model = model self._indices = indices self._velocity = velocity @@ -336,66 +442,24 @@ def alloc(self, num_envs: int) -> np.ndarray: return np.zeros((num_envs, self._indices.size), dtype=np.float32) def __call__(self, buffers, idx: np.ndarray | slice, rows: mtx.SceneData) -> None: - target = rows.dof_vel if self._velocity else rows.dof_pos - target[:, self._indices] = buffers[idx] + if not self._indices.size: + return + # SceneData property projections are read-only copies; channel patches + # reach the sim state only through the explicit dof setters. + if self._velocity: + values = rows.dof_vel + values[:, self._indices] = buffers[idx] + rows.set_dof_vel(values) + else: + values = rows.dof_pos + values[:, self._indices] = buffers[idx] + rows.set_dof_pos(values, self._model) def apply(self, dof_pos, dof_vel, buffers, env_ids) -> None: target = dof_vel if self._velocity else dof_pos target[:, self._indices] = buffers[env_ids] -class _DofComponentPatchOp: - def __init__(self, indices: np.ndarray, *, velocity: bool = False) -> None: - self._indices = indices - self._velocity = velocity - - def apply(self, dof_pos, dof_vel, buffers, env_ids) -> None: - target = dof_vel if self._velocity else dof_pos - values = buffers[env_ids] - for target_index, indices in enumerate(self._indices): - target[:, indices] = values[:, target_index] - - -class _CtrlOp: - """Ctrl targets routed to fixed native actuator columns.""" - - def __init__(self, indices: np.ndarray) -> None: - self.indices = indices - - def alloc(self, num_envs: int) -> np.ndarray: - return np.zeros((num_envs, self.indices.size), dtype=np.float32) - - def __call__(self, buffers, idx: np.ndarray | slice, rows: mtx.SceneData) -> None: - values = buffers[idx] - if not values.shape[1]: - return - if self.indices.size == rows.actuator_ctrls.shape[1]: - rows.actuator_ctrls = values - else: - rows.actuator_ctrls[:, self.indices] = values - - -class _MultiTargetOp: - """Apply one fixed-width property to targets in declared order.""" - - def __init__(self, targets, setter_name: str, width: int, *, contiguous: bool = False) -> None: - self._setters = [getattr(target, setter_name) for target in targets] - self._width = width - self._contiguous = contiguous - - def alloc(self, num_envs: int) -> np.ndarray: - shape = (num_envs, len(self._setters)) if self._width == 1 else (num_envs, len(self._setters), self._width) - return np.zeros(shape, dtype=np.float32) - - def __call__(self, buffers, idx: np.ndarray | slice, rows: mtx.SceneData) -> None: - values = buffers[idx] - for target_index, setter in enumerate(self._setters): - target_values = values[:, target_index] - if self._contiguous: - target_values = np.ascontiguousarray(target_values) - setter(rows, target_values) - - def _named_body(model: mtx.SceneModel, body_name: str): body = model.get_body(body_name) if body is None: diff --git a/motrix_env_motrixsim/tests/test_motrixsim_backend.py b/motrix_env_motrixsim/tests/test_motrixsim_backend.py index ce343636..0d425934 100644 --- a/motrix_env_motrixsim/tests/test_motrixsim_backend.py +++ b/motrix_env_motrixsim/tests/test_motrixsim_backend.py @@ -32,10 +32,11 @@ LinkLinearVelocityQuery, LinkPositionQuery, LinkQuaternionQuery, + SitePositionQuery, ) from motrix_env_core.sim.backend import SimModel from motrix_env_core.sim.registry import create_sim_backend, list_sim_backends -from motrix_env_core.sim.write import BodyJointVelocityWrite, DofVelocityWrite, JointVelocityWrite +from motrix_env_core.sim.write import BodyJointVelocityWrite, DofVelocityWrite, JointVelocityWrite, MocapPoseWrite from motrix_env_motrixsim.compiler import MotrixSimSceneCompiler from motrix_env_motrixsim.runtime import MotrixSimBackend @@ -349,6 +350,92 @@ def test_body_state_reset_is_visible_to_link_queries(): np.testing.assert_allclose(read["angular_velocity"][0], [0.4, 0.5, 0.6]) +def test_mixed_reset_program_applies_native_and_numpy_writes_together(): + import motrix_envs # noqa: F401 + from motrix_env_core import registry + + cfg = registry.make_env_config("dm-humanoid-walk", mode="play") + backend = MotrixSimBackend(cfg.scene, cfg.sim, 2) + body = "torso" + reset = backend.write_compiler.compile( + { + "base_position": BodyPositionWrite((body,)), + "base_rotation": BodyRotationWrite((body,)), + "joint_position": JointPositionWrite(("abdomen_z",)), + "joint_velocity": JointVelocityWrite(("abdomen_z",)), + }, + reset=True, + ) + read = backend.compile_reads( + { + "position": LinkPositionQuery(link=body), + "joint_position": JointPositionQuery(joints=("abdomen_z",)), + "joint_velocity": JointVelocityQuery(joints=("abdomen_z",)), + } + ) + assert reset.buffer("base_position").shape == (2, 1, 3) + reset.buffer("base_position")[1, 0] = [0.5, 0.0, 1.0] + reset.buffer("joint_position")[1] = [0.25] + reset.buffer("joint_velocity")[1] = [-0.75] + + reset.execute(np.asarray([1], dtype=np.int64)) + read.execute() + + # Native reset applies first, then the declared writes land on top. + np.testing.assert_allclose(read["position"][1], [0.5, 0.0, 1.0]) + np.testing.assert_allclose(read["joint_position"][1], [0.25]) + np.testing.assert_allclose(read["joint_velocity"][1], [-0.75]) + # The unselected environment stays at the model initial state. + assert not np.allclose(read["joint_position"][0], [0.25]) + + +def test_dof_channel_writes_without_reset_land_in_sim_state(): + import motrix_envs # noqa: F401 + from motrix_env_core import registry + + cfg = registry.make_env_config("dm-finger-spin", mode="play") + backend = MotrixSimBackend(cfg.scene, cfg.sim, 2) + writes = backend.write_compiler.compile( + {"joint_position": JointPositionWrite(("hinge",)), "joint_velocity": JointVelocityWrite(("hinge",))} + ) + read = backend.compile_reads( + { + "joint_position": JointPositionQuery(joints=("hinge",)), + "joint_velocity": JointVelocityQuery(joints=("hinge",)), + } + ) + writes.buffer("joint_position")[0] = [0.3] + writes.buffer("joint_velocity")[0] = [-0.6] + + writes.execute(np.asarray([0], dtype=np.int64)) + read.execute(np.asarray([0], dtype=np.int64)) + + np.testing.assert_allclose(read["joint_position"][0], [0.3]) + np.testing.assert_allclose(read["joint_velocity"][0], [-0.6]) + np.testing.assert_array_equal(read["joint_velocity"][1], [0.0]) + + +def test_mocap_pose_writes_keep_the_flat_pose_layout_and_reach_the_sim(): + import motrix_envs # noqa: F401 + from motrix_env_core import registry + + cfg = registry.make_env_config("dm-manipulator-bring-ball", mode="play") + backend = MotrixSimBackend(cfg.scene, cfg.sim, 2) + writes = backend.write_compiler.compile({"target": MocapPoseWrite(("target_ball",))}) + read = backend.compile_reads({"pos": SitePositionQuery(site="target_ball")}) + + buffer = writes.buffer("target") + assert buffer.shape == (2, 1, 7) + buffer[1, 0] = [0.11, 0.22, 0.33, 0.0, 0.0, 0.0, 1.0] + + writes.execute(np.asarray([1], dtype=np.int64)) + read.execute(np.asarray([1], dtype=np.int64)) + + np.testing.assert_allclose(read["pos"][1], [0.11, 0.22, 0.33]) + # The other environment keeps its own mocap pose. + np.testing.assert_allclose(read["pos"][0], [0.0, 0.0, 0.0]) + + def test_named_joint_queries_reject_unknown_joints(): import motrix_envs # noqa: F401 from motrix_env_core import registry diff --git a/motrix_env_motrixsim/tests/test_multi_target_writes.py b/motrix_env_motrixsim/tests/test_multi_target_writes.py deleted file mode 100644 index 5f5cc128..00000000 --- a/motrix_env_motrixsim/tests/test_multi_target_writes.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Structured target-axis behavior of fixed-width MotrixSim write ops.""" - -import numpy as np - -from motrix_env_motrixsim.write_compiler import _MultiTargetOp - - -class _Target: - def __init__(self) -> None: - self.values = [] - - def set_scalar(self, rows, values) -> None: - self.values.append((rows, values.copy())) - - def set_vector(self, rows, values) -> None: - self.values.append((rows, values.copy())) - - -def test_multi_target_scalar_op_routes_declared_target_axis() -> None: - first = _Target() - second = _Target() - op = _MultiTargetOp((first, second), "set_scalar", 1) - buffers = op.alloc(3) - assert buffers.shape == (3, 2) - buffers[:] = [[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]] - - op(buffers, np.asarray([2, 0], dtype=np.int64), "rows") - - np.testing.assert_array_equal(first.values[0][1], [3.0, 1.0]) - np.testing.assert_array_equal(second.values[0][1], [30.0, 10.0]) - - -def test_multi_target_vector_op_routes_declared_target_axis() -> None: - first = _Target() - second = _Target() - op = _MultiTargetOp((first, second), "set_vector", 3) - buffers = op.alloc(2) - assert buffers.shape == (2, 2, 3) - buffers[:] = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]] - - op(buffers, slice(None), "rows") - - np.testing.assert_array_equal(first.values[0][1], [[1, 2, 3], [7, 8, 9]]) - np.testing.assert_array_equal(second.values[0][1], [[4, 5, 6], [10, 11, 12]]) diff --git a/motrix_env_motrixsim/tests/test_write_compiler.py b/motrix_env_motrixsim/tests/test_write_compiler.py index 38fe7f90..57a5c569 100644 --- a/motrix_env_motrixsim/tests/test_write_compiler.py +++ b/motrix_env_motrixsim/tests/test_write_compiler.py @@ -77,25 +77,30 @@ def test_reset_program_passes_compile_time_kinematics_flag_to_native_reset() -> assert model.forward_kinematic_rows == [] -def test_reset_program_applies_non_fused_writes_after_native_reset_and_refreshes_once() -> None: +def test_reset_program_folds_ops_into_the_reset_state() -> None: model = _Model() data = _ResetData() - op = _Op() + folded = [] + + class _FoldOp: + def apply(self, dof_pos, dof_vel, buffers, env_ids) -> None: + folded.append((dof_pos.shape, dof_vel.shape, list(env_ids))) + program = _MotrixSimWriteProgram( model, data, - lambda env_ids: env_ids, + lambda env_ids: data, {}, - [(_CompiledWrite(op), {})], + [(_CompiledWrite(_Op(), _FoldOp()), np.zeros((3, 2), dtype=np.float32))], reset=True, refresh_kinematics=True, ) - program.execute() + program.execute(np.asarray([0, 2], dtype=np.int64)) - assert data.reset_calls == [(model, {"forward_kinematic": False})] - assert op.rows == [data] - assert model.forward_kinematic_rows == [data] + # The reset fires once with the folded state; no op runs post-reset. + assert data.reset_calls == [(model, {"forward_kinematic": True})] + assert folded == [((2, 0), (2, 0), [0, 2])] def test_write_program_skips_kinematic_refresh_when_no_op_requires_it() -> None: @@ -107,3 +112,113 @@ def test_write_program_skips_kinematic_refresh_when_no_op_requires_it() -> None: program.execute() assert model.forward_kinematic_rows == [] + + +class _NativeProgram: + def __init__(self, label: str, log: list) -> None: + self._label = label + self._log = log + + def execute(self, data, env_ids) -> None: + self._log.append((self._label, data, None if env_ids is None else env_ids.copy())) + + +class _LoggedOp: + def __init__(self, label: str, log: list) -> None: + self._label = label + self._log = log + + def alloc(self, num_envs: int) -> np.ndarray: + return np.zeros((num_envs, 1), dtype=np.float32) + + def __call__(self, buffers, idx, rows) -> None: + del buffers, idx + self._log.append((self._label, rows)) + + +def _native_case( + log: list, ops: list, *, reset: bool, refresh: bool, lead: int, env_ids=None +) -> _MotrixSimWriteProgram: + model = _Model() + program = _MotrixSimWriteProgram( + model, + _Data(), + lambda env_ids: ("rows", tuple(env_ids)), + {}, + ops, + native=_NativeProgram("native", log), + reset=reset, + refresh_kinematics=refresh, + lead_op_count=lead, + ) + program.execute(env_ids) + return program + + +def test_native_program_runs_lead_ops_then_native_then_tail_ops() -> None: + log = [] + ops = [(_CompiledWrite(_LoggedOp("lead", log)), {}), (_CompiledWrite(_LoggedOp("tail", log)), {})] + + _native_case(log, ops, reset=False, refresh=False, lead=1, env_ids=np.asarray([1], dtype=np.int64)) + + assert [entry[0] for entry in log] == ["lead", "native", "tail"] + assert log[1][2].tolist() == [1] + + +def test_native_reset_program_runs_native_pass_before_numpy_ops() -> None: + log = [] + ops = [(_CompiledWrite(_LoggedOp("op", log)), {})] + + _native_case(log, ops, reset=True, refresh=False, lead=0) + + assert [entry[0] for entry in log] == ["native", "op"] + assert log[0][2] is None + + +def test_native_program_refreshes_kinematics_once_after_every_write() -> None: + log = [] + ops = [(_CompiledWrite(_LoggedOp("op", log)), {})] + model = _Model() + data = _Data() + + program = _MotrixSimWriteProgram( + model, + data, + lambda env_ids: ("rows", tuple(env_ids)), + {}, + ops, + native=_NativeProgram("native", log), + reset=False, + refresh_kinematics=True, + lead_op_count=0, + ) + program.execute() + + assert [entry[0] for entry in log] == ["native", "op"] + assert model.forward_kinematic_rows == [data] + + +def test_native_buffer_lookup_falls_through_to_native_views() -> None: + class _ViewNative: + def __getitem__(self, name: str) -> np.ndarray: + return {"ctrl": np.zeros((2, 3), dtype=np.float32)}[name] + + program = _MotrixSimWriteProgram( + _Model(), + _Data(), + lambda env_ids: env_ids, + {"state": np.zeros((2, 2), dtype=np.float32)}, + [], + native=_ViewNative(), + reset=False, + refresh_kinematics=False, + ) + + assert program.buffer("state").shape == (2, 2) + assert program.buffer("ctrl").shape == (2, 3) + try: + program.buffer("missing") + except KeyError: + pass + else: + raise AssertionError("Unknown buffer name must raise KeyError.") diff --git a/uv.lock b/uv.lock index 842d29a2..279b6bfc 100644 --- a/uv.lock +++ b/uv.lock @@ -1021,7 +1021,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "motrix-env-core", editable = "motrix_env_core" }, - { name = "motrixsim", specifier = "==0.10.1.dev123478" }, + { name = "motrixsim", specifier = "==0.10.1.dev123565" }, { name = "numpy", specifier = ">=1.26" }, ] @@ -1242,27 +1242,27 @@ provides-extras = ["skrl-jax", "skrl-torch", "rslrl"] [[package]] name = "motrixsim" -version = "0.10.1.dev123478" +version = "0.10.1.dev123565" source = { registry = "https://pypi.motphys.com/simple" } dependencies = [ { name = "motrixsim-core" }, ] wheels = [ - { url = "https://pypi.motphys.com/packages/motrixsim-0.10.1.dev123478-py3-none-any.whl", hash = "sha256:53a86b45f9dbd7b3dc89c70e97740d5a582383ae137ecffbd7704696e25788ea" }, + { url = "https://pypi.motphys.com/packages/motrixsim-0.10.1.dev123565-py3-none-any.whl", hash = "sha256:cdb1b0dd98774e54494db438de03e8c371d8bcf05ae0304d160336b7e394122c" }, ] [[package]] name = "motrixsim-core" -version = "0.10.1.dev123478+pro" +version = "0.10.1.dev123565+pro" source = { registry = "https://pypi.motphys.com/simple" } dependencies = [ { name = "absl-py" }, { name = "numpy" }, ] wheels = [ - { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fe585358b7c1dc9a81953e6820873e20e45a80bb226ca4ec287c5cdf72d75109" }, - { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9135536195f7139274ada6507cf1d23e0216ce0627867925bbefa06c706fca7" }, - { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-win_amd64.whl", hash = "sha256:be99be2772c5ef8947c9f4b171b0cc4b68204dcd1fdbe2ab000a963ee46a4252" }, + { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123565+pro-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:44a45535c78c0236cc1c18f4819e725a4a2b43ab4dfb1ee2261db40b8bd2df02" }, + { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123565+pro-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:507bc1bd96ca208d0a3741b45ac86999f1252a864975dc4e2685f6a5b685e5e9" }, + { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123565+pro-cp310-cp310-win_amd64.whl", hash = "sha256:2ed2cdc77f2f3160e85f90bc24932422d1dd405e355f79057bac0b78c15b36ff" }, ] [[package]]