[Question] Action Masking in Discrete Space

[shiven2001]

Hi, I am making an environment where the robot can take the following actions: move 1m left/right/forward/back and 10 degrees rotate left and right. But if the robot approaches the boundary (10x10 grid) then it cannot take an action which will get it outside of the boundary. so if the robot is at (5.0, 0.0, 0.0), then the action which leads to (6.0, 0, 0) is invalid and not chosen... How to implement this?

  from enum import Enum, auto
  from typing import Sequence, Tuple, Type
  import torch
  import logging
  import numpy as np
  from isaaclab.utils import configclass
  from isaaclab.managers import ActionTerm, ActionTermCfg
  from isaaclab.assets import RigidObject
  from isaaclab.envs import DirectRLEnv
  import math
  
  
  class CameraControlMode(Enum):
      DISCRETE = auto()
      CONTINUOUS = auto()
  
  
  class FloatCamDiscreteAction(ActionTerm):
      cfg: "FloatCamDiscreteActionCfg"
      _env: DirectRLEnv
      _asset: RigidObject
  
      def __init__(self, cfg: "FloatCamDiscreteActionCfg", env: DirectRLEnv):
          super().__init__(cfg, env)
          self._asset = env.scene[cfg.asset_name]
          self._setup_action_space()
          self._raw_actions = torch.zeros(
              (self.num_envs, 1), dtype=torch.long, device=self.device
          )
          self._processed_actions = torch.zeros(
              (self.num_envs, 2), device=self.device
          )  # [forward, yaw]
          self._initial_states = self._asset.data.root_state_w.clone()  # (num_envs, 13)
          # Action mapping: [forward, yaw_deg]
          self._action_values = torch.tensor(
              [
                  [1.0, 0.0, 0.0],  # 0: +1 x
                  [-1.0, 0.0, 0.0],  # 1: -1 x
                  [0.0, 1.0, 0.0],  # 2: +1 y
                  [0.0, -1.0, 0.0],  # 3: -1 y
                  [0.0, 0.0, 10.0],  # 4: +10 deg z rot
                  [0.0, 0.0, -10.0],  # 5: -10 deg z rot
                  [0.0, 0.0, 0.0],  # 6: no-op
              ],
              device=self.device,
          )  # (7, 3)
  
      @property
      def action_dim(self) -> int:
          return 1  # Single discrete action index
  
      @property
      def raw_actions(self) -> torch.Tensor:
          return self._raw_actions
  
      @property
      def processed_actions(self) -> torch.Tensor:
          return self._processed_actions
  
      def process_actions(self, actions: torch.Tensor):
          self._raw_actions = actions  # (num_envs,)
          self._processed_actions = self._action_values[
              actions.squeeze(-1)
          ]  # (num_envs, 3)
  
      def apply_actions(self):
          actions = self._processed_actions  # (num_envs, 3)
          dx = actions[:, 0]  # (num_envs,)
          dy = actions[:, 1]  # (num_envs,)
          yaw_deg = actions[:, 2]  # (num_envs,)
          root_states = self._asset.data.root_state_w  # (num_envs, 13)
          positions = root_states[:, :3]  # (num_envs, 3)
          quats = root_states[:, 3:7]  # (num_envs, 4)
          # Move in x and y in local frame
          if torch.any((dx != 0) | (dy != 0)):
              # Local x and y unit vectors
              local_x = torch.tensor(
                  [1, 0, 0], device=self.device, dtype=positions.dtype
              ).expand(positions.shape[0], 3)
              local_y = torch.tensor(
                  [0, 1, 0], device=self.device, dtype=positions.dtype
              ).expand(positions.shape[0], 3)
  
              def quat_rotate(q, v):
                  w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
                  q_vec = torch.stack([x, y, z], dim=1)
                  t = 2 * torch.cross(q_vec, v, dim=1)
                  v_rot = v + w.unsqueeze(1) * t + torch.cross(q_vec, t, dim=1)
                  return v_rot
  
              world_x = quat_rotate(quats, local_x)
              world_y = quat_rotate(quats, local_y)
              positions = (
                  positions + world_x * dx.unsqueeze(1) + world_y * dy.unsqueeze(1)
              )
          # Rotate about z (yaw)
          if torch.any(yaw_deg != 0):
              yaw_rad = yaw_deg * math.pi / 180.0
              cos = torch.cos(yaw_rad / 2)
              sin = torch.sin(yaw_rad / 2)
              zeros = torch.zeros_like(cos)
              yaw_quat = torch.stack([cos, zeros, zeros, sin], dim=1)
  
              def quat_mult(q1, q2):
                  w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3]
                  w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3]
                  w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
                  x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
                  y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
                  z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
                  return torch.stack([w, x, y, z], dim=1)
  
              quats = quat_mult(yaw_quat, quats)
          # Velocities zero
          linear_velocities = torch.zeros((self.num_envs, 3), device=self.device)
          angular_velocities = torch.zeros((self.num_envs, 3), device=self.device)
          root_state = torch.cat(
              [positions, quats, linear_velocities, angular_velocities], dim=-1
          )  # (num_envs, 13)
          env_ids = torch.arange(self.num_envs, device=self.device)
          self._asset.write_root_state_to_sim(root_state, env_ids=env_ids)
  
      def reset(self, env_ids: Sequence[int] | None = None):
          if env_ids is None:
              env_ids = torch.arange(self.num_envs, device=self.device).tolist()
          else:
              env_ids = list(env_ids)
          self._raw_actions[env_ids] = 0
          self._processed_actions[env_ids] = 0.0
  
      def _setup_action_space(self):
          pass  # No-op for now
  
  
  @configclass
  class FloatCamDiscreteActionCfg(ActionTermCfg):
      class_type: Type = FloatCamDiscreteAction
      asset_name: str = "float_cam"
      control_mode: CameraControlMode = CameraControlMode.DISCRETE

[RandomOakForest]

Thank you for posting this. I'll move this post to our Discussions. Here are a few things to consider.

1. Check Next State Validity Before Action
   Before applying the action in your env.step() (or _apply_action in Isaac Lab direct workflow), calculate the next potential state that would result from each action. Compare this against your grid boundaries.

2. Mask or Ignore Invalid Actions

   • Masking (recommended for RL):
     Create an action mask indicating which actions are currently valid, given the robot's state. Many RL libraries support passing such a mask, so the agent can only sample or select valid actions.
   • In step/reward (Gym/Isaac style):
     In your step method, if an invalid action is chosen (i.e., would go out of bounds), simply do not move the robot and (optionally) provide negative reward or no reward for that step. This discourages invalid moves during policy learning, but does not prevent the agent from selecting them, so masking is preferred.