Move the Direct observation kernels into the family kernel module

Relocates out_of_reach/full_obs/reduced_obs kernels from
reorient_direct_env.py to reorient_kernels.py and lifts the stray
mid-file imports to the top of the module.

Author: hujc7

source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py

@@ -25,162 +25,13 @@
     ReorientRewardBuffers,
     direct_reorient_reward,
     ema_actuation_kernel,
-    fingertip_pos_col,
-    fingertip_quat_col,
-    fingertip_vel_col,
+    full_obs_kernel,
+    out_of_reach_kernel,
+    reduced_obs_kernel,
     reorient_success_kernel,
 )
-from isaaclab_tasks.core.utils import randomize_rotation
-
-
-@wp.kernel
-def _out_of_reach_kernel(
-    object_pos_w: wp.array(dtype=wp.vec3f),
-    env_origins: wp.array(dtype=wp.vec3f),
-    target_pos_e: wp.array(dtype=wp.vec3f),
-    fall_distance: float,
-    out_of_reach: wp.array(dtype=wp.bool),
-):
-    i = wp.tid()
-    out_of_reach[i] = wp.length(object_pos_w[i] - env_origins[i] - target_pos_e[i]) >= fall_distance
-
-
-@wp.kernel
-def _full_obs_kernel(
-    joint_pos: wp.array2d(dtype=wp.float32),
-    joint_vel: wp.array2d(dtype=wp.float32),
-    lower: wp.array2d(dtype=wp.float32),
-    upper: wp.array2d(dtype=wp.float32),
-    vel_scale: float,
-    object_pos_w: wp.array(dtype=wp.vec3f),
-    env_origins: wp.array(dtype=wp.vec3f),
-    object_quat: wp.array(dtype=wp.quatf),
-    object_lin_vel: wp.array(dtype=wp.vec3f),
-    object_ang_vel: wp.array(dtype=wp.vec3f),
-    in_hand_pos_e: wp.array(dtype=wp.vec3f),
-    goal_quat: wp.array(dtype=wp.quatf),
-    body_pos_w: wp.array2d(dtype=wp.vec3f),
-    body_quat_w: wp.array2d(dtype=wp.quatf),
-    body_vel_w: wp.array2d(dtype=wp.spatial_vectorf),
-    finger_ids: wp.array(dtype=wp.int32),
-    force: wp.array2d(dtype=wp.vec3f),
-    torque: wp.array2d(dtype=wp.vec3f),
-    wrench_ids: wp.array(dtype=wp.int32),
-    force_scale: float,
-    with_forces: int,
-    actions: wp.array2d(dtype=wp.float32),
-    out: wp.array2d(dtype=wp.float32),
-):
-    """Direct full observation / full state, matching the torch concatenation order.
-
-    Launched over ``(num_envs, obs_dim)``: each thread walks a branch ladder over the
-    segment boundaries and writes one output column, so warps (32 consecutive columns)
-    stay branch-uniform except at segment boundaries.
-    """
-    i, j = wp.tid()
-    num_joints = joint_pos.shape[1]
-    num_fingers = finger_ids.shape[0]
-    # segment boundaries, in column order
-    end_joint = 2 * num_joints
-    end_object = end_joint + 13
-    end_goal = end_object + 11
-    end_tip_pos = end_goal + 3 * num_fingers
-    end_tip_quat = end_tip_pos + 4 * num_fingers
-    end_tip_vel = end_tip_quat + 6 * num_fingers
-    end_wrench = end_tip_vel
-    if with_forces != 0:
-        end_wrench += 6 * num_fingers
-    # hand: normalized DOF positions, scaled DOF velocities
-    if j < num_joints:
-        out[i, j] = 2.0 * (joint_pos[i, j] - lower[i, j]) / (upper[i, j] - lower[i, j]) - 1.0
-    elif j < end_joint:
-        out[i, j] = vel_scale * joint_vel[i, j - num_joints]
-    # object pose and velocities (environment frame position)
-    elif j < end_object:
-        k = j - end_joint
-        if k < 3:
-            p = object_pos_w[i] - env_origins[i]
-            out[i, j] = p[k]
-        elif k < 7:
-            out[i, j] = object_quat[i][k - 3]
-        elif k < 10:
-            out[i, j] = object_lin_vel[i][k - 7]
-        else:
-            out[i, j] = vel_scale * object_ang_vel[i][k - 10]
-    # goal: in-hand anchor, goal rotation, and the goal-to-object rotation difference
-    elif j < end_goal:
-        k = j - end_object
-        if k < 3:
-            out[i, j] = in_hand_pos_e[i][k]
-        elif k < 7:
-            out[i, j] = goal_quat[i][k - 3]
-        else:
-            # quat_inverse == conjugate for these unit quaternions, matching
-            # isaaclab.utils.math.quat_mul/quat_conjugate semantics
-            qe = object_quat[i] * wp.quat_inverse(goal_quat[i])
-            out[i, j] = qe[k - 7]
-    # fingertips: environment-frame positions, rotations, spatial velocities
-    elif j < end_tip_pos:
-        out[i, j] = fingertip_pos_col(body_pos_w, env_origins, finger_ids, i, j - end_goal)
-    elif j < end_tip_quat:
-        out[i, j] = fingertip_quat_col(body_quat_w, finger_ids, i, j - end_tip_pos)
-    elif j < end_tip_vel:
-        out[i, j] = fingertip_vel_col(body_vel_w, finger_ids, i, j - end_tip_quat)
-    # fingertip force/torque sensors (full state only; absent when with_forces == 0)
-    elif j < end_wrench:
-        if with_forces == 1:
-            k = j - end_tip_vel
-            c = k % 6
-            if c < 3:
-                out[i, j] = force_scale * force[i, wrench_ids[k // 6]][c]
-            else:
-                out[i, j] = force_scale * torque[i, wrench_ids[k // 6]][c - 3]
-        else:
-            # full state requested but the sensor has no data yet: zero block
-            out[i, j] = 0.0
-    # actions
-    else:
-        out[i, j] = actions[i, j - end_wrench]
-
-
-@wp.kernel
-def _reduced_obs_kernel(
-    body_pos_w: wp.array2d(dtype=wp.vec3f),
-    env_origins: wp.array(dtype=wp.vec3f),
-    finger_ids: wp.array(dtype=wp.int32),
-    object_pos_w: wp.array(dtype=wp.vec3f),
-    object_quat: wp.array(dtype=wp.quatf),
-    goal_quat: wp.array(dtype=wp.quatf),
-    actions: wp.array2d(dtype=wp.float32),
-    out: wp.array2d(dtype=wp.float32),
-):
-    """Direct reduced (OpenAI) observation, matching the torch concatenation order."""
-    i = wp.tid()
-    num_fingers = finger_ids.shape[0]
-    for f in range(num_fingers):
-        fp = body_pos_w[i, finger_ids[f]] - env_origins[i]
-        out[i, 3 * f + 0] = fp[0]
-        out[i, 3 * f + 1] = fp[1]
-        out[i, 3 * f + 2] = fp[2]
-    idx = 3 * num_fingers
-    p = object_pos_w[i] - env_origins[i]
-    # quat_inverse == conjugate for these unit quaternions, matching
-    # isaaclab.utils.math.quat_mul/quat_conjugate semantics
-    qe = object_quat[i] * wp.quat_inverse(goal_quat[i])
-    out[i, idx + 0] = p[0]
-    out[i, idx + 1] = p[1]
-    out[i, idx + 2] = p[2]
-    out[i, idx + 3] = qe[0]
-    out[i, idx + 4] = qe[1]
-    out[i, idx + 5] = qe[2]
-    out[i, idx + 6] = qe[3]
-    idx += 7
-    for a in range(actions.shape[1]):
-        out[i, idx + a] = actions[i, a]
-
-
 from isaaclab_tasks.core.reorient.reorient_task_base import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET
-from isaaclab_tasks.core.utils import EpisodeErrorRecorder, sample_joint_positions_within_limits
+from isaaclab_tasks.core.utils import EpisodeErrorRecorder, randomize_rotation, sample_joint_positions_within_limits
 
 if TYPE_CHECKING:
     from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg
@@ -422,7 +273,7 @@ def _get_rewards(self) -> torch.Tensor:
     def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
         # reset when cube has fallen
         wp.launch(
-            _out_of_reach_kernel,
+            out_of_reach_kernel,
             dim=self.num_envs,
             inputs=[
                 self.object.data.root_pos_w.warp,
@@ -545,7 +396,7 @@ def _compute_intermediate_values(self):
     def compute_reduced_observations(self):
         # Per https://arxiv.org/pdf/1808.00177.pdf Table 2
         wp.launch(
-            _reduced_obs_kernel,
+            reduced_obs_kernel,
             dim=self.num_envs,
             inputs=[
                 self.hand.data.body_pos_w.warp,
@@ -564,7 +415,7 @@ def compute_reduced_observations(self):
     def compute_full_observations(self):
         force_wp, torque_wp = self._dummy_wrench_wp, self._dummy_wrench_wp
         wp.launch(
-            _full_obs_kernel,
+            full_obs_kernel,
             dim=(self.num_envs, self._policy_obs_buf_wp.shape[1]),
             inputs=[
                 self.hand.data.joint_pos.warp,
@@ -603,7 +454,7 @@ def compute_full_state(self):
             if force_data is not None and torque_data is not None:
                 force_wp, torque_wp, with_forces = force_data.warp, torque_data.warp, 1
         wp.launch(
-            _full_obs_kernel,
+            full_obs_kernel,
             dim=(self.num_envs, self._state_obs_buf_wp.shape[1]),
             inputs=[
                 self.hand.data.joint_pos.warp,

source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_kernels.py

@@ -423,3 +423,149 @@ def compute_cube_keypoints(
         device=wp.device_from_torch(pose.device),
     )
     return out
+
+
+@wp.kernel
+def out_of_reach_kernel(
+    object_pos_w: wp.array(dtype=wp.vec3f),
+    env_origins: wp.array(dtype=wp.vec3f),
+    target_pos_e: wp.array(dtype=wp.vec3f),
+    fall_distance: float,
+    out_of_reach: wp.array(dtype=wp.bool),
+):
+    i = wp.tid()
+    out_of_reach[i] = wp.length(object_pos_w[i] - env_origins[i] - target_pos_e[i]) >= fall_distance
+
+
+@wp.kernel
+def full_obs_kernel(
+    joint_pos: wp.array2d(dtype=wp.float32),
+    joint_vel: wp.array2d(dtype=wp.float32),
+    lower: wp.array2d(dtype=wp.float32),
+    upper: wp.array2d(dtype=wp.float32),
+    vel_scale: float,
+    object_pos_w: wp.array(dtype=wp.vec3f),
+    env_origins: wp.array(dtype=wp.vec3f),
+    object_quat: wp.array(dtype=wp.quatf),
+    object_lin_vel: wp.array(dtype=wp.vec3f),
+    object_ang_vel: wp.array(dtype=wp.vec3f),
+    in_hand_pos_e: wp.array(dtype=wp.vec3f),
+    goal_quat: wp.array(dtype=wp.quatf),
+    body_pos_w: wp.array2d(dtype=wp.vec3f),
+    body_quat_w: wp.array2d(dtype=wp.quatf),
+    body_vel_w: wp.array2d(dtype=wp.spatial_vectorf),
+    finger_ids: wp.array(dtype=wp.int32),
+    force: wp.array2d(dtype=wp.vec3f),
+    torque: wp.array2d(dtype=wp.vec3f),
+    wrench_ids: wp.array(dtype=wp.int32),
+    force_scale: float,
+    with_forces: int,
+    actions: wp.array2d(dtype=wp.float32),
+    out: wp.array2d(dtype=wp.float32),
+):
+    """Direct full observation / full state, matching the torch concatenation order.
+
+    Launched over ``(num_envs, obs_dim)``: each thread walks a branch ladder over the
+    segment boundaries and writes one output column, so warps (32 consecutive columns)
+    stay branch-uniform except at segment boundaries.
+    """
+    i, j = wp.tid()
+    num_joints = joint_pos.shape[1]
+    num_fingers = finger_ids.shape[0]
+    # segment boundaries, in column order
+    end_joint = 2 * num_joints
+    end_object = end_joint + 13
+    end_goal = end_object + 11
+    end_tip_pos = end_goal + 3 * num_fingers
+    end_tip_quat = end_tip_pos + 4 * num_fingers
+    end_tip_vel = end_tip_quat + 6 * num_fingers
+    end_wrench = end_tip_vel
+    if with_forces != 0:
+        end_wrench += 6 * num_fingers
+    # hand: normalized DOF positions, scaled DOF velocities
+    if j < num_joints:
+        out[i, j] = 2.0 * (joint_pos[i, j] - lower[i, j]) / (upper[i, j] - lower[i, j]) - 1.0
+    elif j < end_joint:
+        out[i, j] = vel_scale * joint_vel[i, j - num_joints]
+    # object pose and velocities (environment frame position)
+    elif j < end_object:
+        k = j - end_joint
+        if k < 3:
+            p = object_pos_w[i] - env_origins[i]
+            out[i, j] = p[k]
+        elif k < 7:
+            out[i, j] = object_quat[i][k - 3]
+        elif k < 10:
+            out[i, j] = object_lin_vel[i][k - 7]
+        else:
+            out[i, j] = vel_scale * object_ang_vel[i][k - 10]
+    # goal: in-hand anchor, goal rotation, and the goal-to-object rotation difference
+    elif j < end_goal:
+        k = j - end_object
+        if k < 3:
+            out[i, j] = in_hand_pos_e[i][k]
+        elif k < 7:
+            out[i, j] = goal_quat[i][k - 3]
+        else:
+            # quat_inverse == conjugate for these unit quaternions, matching
+            # isaaclab.utils.math.quat_mul/quat_conjugate semantics
+            qe = object_quat[i] * wp.quat_inverse(goal_quat[i])
+            out[i, j] = qe[k - 7]
+    # fingertips: environment-frame positions, rotations, spatial velocities
+    elif j < end_tip_pos:
+        out[i, j] = fingertip_pos_col(body_pos_w, env_origins, finger_ids, i, j - end_goal)
+    elif j < end_tip_quat:
+        out[i, j] = fingertip_quat_col(body_quat_w, finger_ids, i, j - end_tip_pos)
+    elif j < end_tip_vel:
+        out[i, j] = fingertip_vel_col(body_vel_w, finger_ids, i, j - end_tip_quat)
+    # fingertip force/torque sensors (full state only; absent when with_forces == 0)
+    elif j < end_wrench:
+        if with_forces == 1:
+            k = j - end_tip_vel
+            c = k % 6
+            if c < 3:
+                out[i, j] = force_scale * force[i, wrench_ids[k // 6]][c]
+            else:
+                out[i, j] = force_scale * torque[i, wrench_ids[k // 6]][c - 3]
+        else:
+            # full state requested but the sensor has no data yet: zero block
+            out[i, j] = 0.0
+    # actions
+    else:
+        out[i, j] = actions[i, j - end_wrench]
+
+
+@wp.kernel
+def reduced_obs_kernel(
+    body_pos_w: wp.array2d(dtype=wp.vec3f),
+    env_origins: wp.array(dtype=wp.vec3f),
+    finger_ids: wp.array(dtype=wp.int32),
+    object_pos_w: wp.array(dtype=wp.vec3f),
+    object_quat: wp.array(dtype=wp.quatf),
+    goal_quat: wp.array(dtype=wp.quatf),
+    actions: wp.array2d(dtype=wp.float32),
+    out: wp.array2d(dtype=wp.float32),
+):
+    """Direct reduced (OpenAI) observation, matching the torch concatenation order."""
+    i = wp.tid()
+    num_fingers = finger_ids.shape[0]
+    for f in range(num_fingers):
+        fp = body_pos_w[i, finger_ids[f]] - env_origins[i]
+        out[i, 3 * f + 0] = fp[0]
+        out[i, 3 * f + 1] = fp[1]
+        out[i, 3 * f + 2] = fp[2]
+    idx = 3 * num_fingers
+    p = object_pos_w[i] - env_origins[i]
+    # quat_inverse == conjugate for these unit quaternions, matching
+    # isaaclab.utils.math.quat_mul/quat_conjugate semantics
+    qe = object_quat[i] * wp.quat_inverse(goal_quat[i])
+    out[i, idx + 0] = p[0]
+    out[i, idx + 1] = p[1]
+    out[i, idx + 2] = p[2]
+    out[i, idx + 3] = qe[0]
+    out[i, idx + 4] = qe[1]
+    out[i, idx + 5] = qe[2]
+    out[i, idx + 6] = qe[3]
+    idx += 7
+    for a in range(actions.shape[1]):
+        out[i, idx + a] = actions[i, a]