diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index 5837114b..d6921ae0 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -21,6 +21,7 @@ wp_gather_indexed_cols_2d, wp_gather_indexed_rows_vec3f, wp_gather_indexed_rows_quatf, + wp_masked_copy_rows_2d, reset_data_keyframe, ) @@ -85,6 +86,53 @@ def reset(self) -> None: self.mj_model, self.mjw_model, self.mjw_data, self._neutral_keyframe_id ) + def set_state( + self, + qpos: Float[np.ndarray | wp.array, "n_worlds nq"] | None = None, + qvel: Float[np.ndarray | wp.array, "n_worlds nv"] | None = None, + worlds_mask: wp.array | np.ndarray | None = None, + ) -> None: + """Directly set generalized positions and/or velocities for parallel worlds. + + Unlike `reset`, which restores the single neutral keyframe to every (masked) + world, this writes a possibly different `qpos`/`qvel` per world. It is the + building block for reference-state initialization, where each world starts an + episode from a different reference pose. + + Args: + qpos: Optional generalized positions, shape ``(n_worlds, nq)``. Accepts + numpy or Warp arrays. If None, positions are left unchanged. + qvel: Optional generalized velocities, shape ``(n_worlds, nv)``. Accepts + numpy or Warp arrays. If None, velocities are left unchanged. + worlds_mask: Optional per-world boolean mask, shape ``(n_worlds,)``. Only + masked worlds are written; the rest keep their current state. Defaults + to all worlds. + + !!! warning + + Like `reset`, this does not update derived kinematic quantities (`xpos`, + `xquat`, `site_xpos`, velocities, ...). Call `forward` (or `step`) before + reading them, otherwise getters return stale pre-write values. + """ + if qpos is None and qvel is None: + return + if worlds_mask is None: + mask = wp.ones(self.n_worlds, dtype=wp.bool) + elif isinstance(worlds_mask, wp.array): + mask = worlds_mask + else: + mask = wp.array(worlds_mask, dtype=wp.bool) + for src, dst in ((qpos, self.mjw_data.qpos), (qvel, self.mjw_data.qvel)): + if src is None: + continue + if not isinstance(src, wp.array): + src = wp.array(src, dtype=wp.float32) + wp.launch( + wp_masked_copy_rows_2d, + dim=dst.shape, + inputs=[src, dst, mask], + ) + @override def get_joint_angles( self, fly_name: str, dst: wp.array | None = None @@ -135,6 +183,74 @@ def get_joint_velocities( ) return dst + def get_base_linear_velocity( + self, fly_name: str, dst: wp.array | None = None + ) -> Float[wp.array, "n_worlds 3"]: + """Get the linear velocity of the fly's base (free joint) for all worlds. + + The base velocity is not part of ``get_joint_velocities`` (which covers only the + fly's articulated DOFs). It is only defined for a fly attached to the world by a + free joint (the usual walking setup); tethered flies have none. + + Args: + fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. + + Returns: + Warp array of shape ``(n_worlds, 3)`` in mm/s, in the **world** frame. + + Raises: + KeyError: If the fly has no free (base) joint. + """ + indices = self._require_base_veladrs(fly_name, self._wp_base_linveladrs_by_fly) + if dst is None: + dst = wp.zeros((self.n_worlds, 3), dtype=wp.float32) + wp.launch( + wp_gather_indexed_cols_2d, + dim=(self.n_worlds, 3), + inputs=[self.mjw_data.qvel, dst, indices], + ) + return dst + + def get_base_angular_velocity( + self, fly_name: str, dst: wp.array | None = None + ) -> Float[wp.array, "n_worlds 3"]: + """Get the angular velocity of the fly's base (free joint) for all worlds. + + See ``get_base_linear_velocity`` for when the base velocity is defined. + + Args: + fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. + + Returns: + Warp array of shape ``(n_worlds, 3)`` in rad/s, in the **root-body** frame + (MuJoCo's native free-joint angular-velocity convention). + + Raises: + KeyError: If the fly has no free (base) joint. + """ + indices = self._require_base_veladrs(fly_name, self._wp_base_angveladrs_by_fly) + if dst is None: + dst = wp.zeros((self.n_worlds, 3), dtype=wp.float32) + wp.launch( + wp_gather_indexed_cols_2d, + dim=(self.n_worlds, 3), + inputs=[self.mjw_data.qvel, dst, indices], + ) + return dst + + @staticmethod + def _require_base_veladrs(fly_name: str, adrs_by_fly: dict) -> wp.array: + if fly_name not in adrs_by_fly: + raise KeyError( + f"Fly '{fly_name}' has no free (base) joint, so its base velocity is " + "undefined (is it tethered?)." + ) + return adrs_by_fly[fly_name] + @override def get_body_positions( self, fly_name: str, dst: wp.array | None = None @@ -294,6 +410,16 @@ def step(self) -> None: """Advance all parallel worlds by one timestep on the GPU.""" mjw.step(self.mjw_model, self.mjw_data) + def forward(self) -> None: + """Recompute derived quantities from the current state, without integrating. + + The batched analogue of MuJoCo's ``mj_forward``: refreshes ``xpos``, ``xquat``, + velocities, and other derived fields from ``qpos``/``qvel`` without advancing + time. Call this after writing state directly (e.g. via ``set_state`` or + ``reset``) so subsequent getters return up-to-date values. + """ + mjw.forward(self.mjw_model, self.mjw_data) + @override def set_renderer( self, @@ -416,6 +542,23 @@ def _map_internal_qposqveladrs(self) -> None: k: wp.array(v, dtype=wp.int32) for k, v in self._intern_qveladrs_by_fly.items() } + # Base ("free") joint velocity addresses, used by get_base_*_velocity. The world + # attaches each walking fly's root body to the world with a free joint named + # after the fly; its 6 velocity DOFs are laid out as [linear xyz, angular xyz]. + # Flies without a free joint (e.g. tethered) are simply left out of these maps. + self._wp_base_linveladrs_by_fly = {} + self._wp_base_angveladrs_by_fly = {} + for fly_name in self._intern_qveladrs_by_fly: + jid = mj.mj_name2id(self.mj_model, mj.mjtObj.mjOBJ_JOINT, fly_name) + if jid < 0 or self.mj_model.jnt_type[jid] != mj.mjtJoint.mjJNT_FREE: + continue + vadr = int(self.mj_model.jnt_dofadr[jid]) + self._wp_base_linveladrs_by_fly[fly_name] = wp.array( + [vadr, vadr + 1, vadr + 2], dtype=wp.int32 + ) + self._wp_base_angveladrs_by_fly[fly_name] = wp.array( + [vadr + 3, vadr + 4, vadr + 5], dtype=wp.int32 + ) @override def _map_internal_actuator_ids(self) -> None: diff --git a/src/flygym/warp/utils.py b/src/flygym/warp/utils.py index 80c409d0..13662ba1 100644 --- a/src/flygym/warp/utils.py +++ b/src/flygym/warp/utils.py @@ -119,6 +119,29 @@ def wp_gather_indexed_cols_2d( dst[i, k] = src[i, cols[k]] +@wp.kernel +def wp_masked_copy_rows_2d( + src: wp.array2d[float], dst: wp.array2d[float], world_mask: wp.array[bool] +): + """Copy whole rows (dim 0, indexed by world) where a per-world mask is set. + + This kernel is to be launched with a 2D launch configuration of + `(n_worlds, n_cols)`. Unlike a keyframe reset, `src` may hold a different row per + world, so this is the building block for per-world reference-state initialization. + + Args: + src (wp.array of shape (n_worlds, n_cols), type float32): + Source array; rows for un-masked worlds are ignored. + dst (wp.array of shape (n_worlds, n_cols), type float32): + Destination array, written in place for masked worlds only. + world_mask (wp.array of shape (n_worlds,), type bool): + Per-world mask; a row is copied only where its entry is True. + """ + i, k = wp.tid() + if world_mask[i]: + dst[i, k] = src[i, k] + + @wp.kernel def unpack_rgb_kernel_selected_worlds_and_cameras( # In: