Final cleanup

Author: hujc7

source/isaaclab_experimental/isaaclab_experimental/envs/MANAGER_BASED_WARP_MIGRATION_PLAN.md

@@ -919,7 +919,7 @@ def _reset_idx(self, env_ids: Sequence[int]):
         info = self.observation_manager.reset(env_mask=env_mask)
         self.extras["log"].update(info)
         # -- action manager
-        info = self.action_manager.reset(mask=env_mask)
+        info = self.action_manager.reset(env_mask=env_mask)
         self.extras["log"].update(info)
         # -- event manager
         info = self.event_manager.reset(env_mask=env_mask)

source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py

@@ -636,12 +636,12 @@ def _reset_idx(
             action_info = self._manager_call_switch.call_stage(
                 stage="ActionManager_reset",
                 stable_calls=[{"fn": self.action_manager.reset, "kwargs": {"env_ids": env_ids}}],
-                warp_calls=[{"fn": self.action_manager.reset, "kwargs": {"mask": env_mask}}],
+                warp_calls=[{"fn": self.action_manager.reset, "kwargs": {"env_mask": env_mask}}],
             )
             reward_info = self._manager_call_switch.call_stage(
                 stage="RewardManager_reset",
                 stable_calls=[{"fn": self.reward_manager.reset, "kwargs": {"env_ids": env_ids}}],
-                warp_calls=[{"fn": self.reward_manager.reset, "kwargs": {"mask": env_mask}}],
+                warp_calls=[{"fn": self.reward_manager.reset, "kwargs": {"env_mask": env_mask}}],
             )
             if self._manager_call_switch.get_mode_for_manager("ObservationManager") == ManagerCallMode.WARP_CAPTURED:
                 obs_info = {}

source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py

@@ -148,7 +148,7 @@ def __init__(self, cfg: actions_cfg.JointActionCfg, env: ManagerBasedEnv) -> Non
 
         # parse offset
         if isinstance(cfg.offset, (float, int)):
-            self._offset = wp.array([float(cfg.offset)], dtype=wp.float32, device=self.device)
+            self._offset = wp.array([float(cfg.offset)] * self.action_dim, dtype=wp.float32, device=self.device)
         elif isinstance(cfg.offset, dict):
             offset_per_joint = [0.0] * self.action_dim
             # resolve the dictionary config
@@ -253,15 +253,15 @@ def process_actions(self, actions: wp.array, action_offset: int = 0):
             device=self.device,
         )
 
-    def reset(self, mask: wp.array | None = None) -> None:
+    def reset(self, env_mask: wp.array | None = None) -> None:
         """Resets the action term (mask-based)."""
-        if mask is None:
+        if env_mask is None:
             self._raw_actions.fill_(0.0)
             return
         wp.launch(
             kernel=_zero_masked_2d,
             dim=(self.num_envs, self.action_dim),
-            inputs=[mask, self._raw_actions],
+            inputs=[env_mask, self._raw_actions],
             device=self.device,
         )