Action Masking and Precheck - SMDPfier
Description: SMDP-level behavior for Gymnasium environments

Skip to content

# Masking and Precheck

SMDPfier provides sophisticated action masking and precheck validation to handle invalid options gracefully and ensure robust option execution.

## Action Masking

Action masking restricts which options are available based on the current environment state. This is essential for handling state-dependent action validity in complex environments.

### Key Concepts

- Availability Function: Determines which option indices are valid
- Action Mask: Binary array indicating available options
- Index Interface Only: Masking only works with `action_interface="index"`
- Dynamic Evaluation: Mask is recomputed every step

### Basic Action Masking

    import gymnasium as gym
    from smdpfier import SMDPfier, Option
    from smdpfier.defaults import ConstantOptionDuration

    # Define options for CartPole
    options = [
        Option([0, 0], "strong-left"),      # Index 0
        Option([1, 1], "strong-right"),     # Index 1
        Option([0, 1], "left-right"),       # Index 2
        Option([1, 0], "right-left"),       # Index 3
    ]

    def cart_availability(obs):
        """Restrict options based on cart position and velocity."""
        position, velocity = obs[0], obs[1]

        available = []

        # Always allow balanced options
        available.extend([2, 3])  # left-right, right-left

        # Restrict strong movements based on position
        if position < 0.3:    # Cart not too far right
            available.append(1)   # Allow strong-right
        if position > -0.3:   # Cart not too far left
            available.append(0)   # Allow strong-left

        return available

    env = SMDPfier(
        gym.make("CartPole-v1"),
        options_provider=options,
        duration_fn=ConstantOptionDuration(5),
        action_interface="index",
        max_options=4,
        availability_fn=cart_availability
    )

    obs, info = env.reset()
    print(f"Available options: {info['smdp']['action_mask']}")
    # Might show: [1, 1, 1, 1] (all available) or [0, 1, 1, 1] (strong-left masked)

### Action Mask Structure

The action mask is a binary list where `1` means available and `0` means masked:

    action_mask = [1, 0, 1, 0]  # Options 0 and 2 available, 1 and 3 masked

Usage in RL algorithms:

    obs, info = env.reset()
    action_mask = info["smdp"]["action_mask"]

    # In your RL algorithm:
    if action_mask is not None:
        # Mask invalid actions (set their Q-values to -inf)
        masked_q_values = q_values.copy()
        masked_q_values[action_mask == 0] = -float('inf')
        action = np.argmax(masked_q_values)
    else:
        action = np.argmax(q_values)

### Complex Masking Examples

#### Environment-Specific Masking

    def taxi_availability(obs):
        """Taxi-v3 environment masking."""
        # Decode Taxi state
        taxi_row, taxi_col, passenger_loc, destination = env.unwrapped.decode(obs)

        available = []

        # Movement actions (always available)
        available.extend([0, 1, 2, 3])  # south, north, east, west

        # Pickup action (only if passenger at taxi location)
        if passenger_loc < 4:  # Passenger not in taxi
            passenger_coords = env.unwrapped.locs[passenger_loc]
            if (taxi_row, taxi_col) == passenger_coords:
                available.append(4)  # Allow pickup

        # Dropoff action (only if passenger in taxi at destination)
        if passenger_loc == 4:  # Passenger in taxi
            destination_coords = env.unwrapped.locs[destination]
            if (taxi_row, taxi_col) == destination_coords:
                available.append(5)  # Allow dropoff

        return available

    # Taxi options
    taxi_options = [
        Option([0], "south"),      # Index 0
        Option([1], "north"),      # Index 1
        Option([2], "east"),       # Index 2
        Option([3], "west"),       # Index 3
        Option([4], "pickup"),     # Index 4
        Option([5], "dropoff"),    # Index 5
    ]

    taxi_env = SMDPfier(
        gym.make("Taxi-v3"),
        options_provider=taxi_options,
        duration_fn=ConstantOptionDuration(1),
        action_interface="index",
        max_options=6,
        availability_fn=taxi_availability
    )

#### State-Dependent Option Length

    def adaptive_availability(obs):
        """Allow different option lengths based on state."""
        velocity = abs(obs[1])  # Cart velocity

        if velocity > 0.5:  # High velocity - need quick corrections
            return [0, 1]   # Only single-action options
        else:  # Low velocity - can use longer sequences
            return [0, 1, 2, 3, 4]  # All options available

    options = [
        Option([0], "quick-left"),        # Index 0 - quick
        Option([1], "quick-right"),       # Index 1 - quick
        Option([0, 0], "double-left"),    # Index 2 - longer
        Option([1, 1], "double-right"),   # Index 3 - longer
        Option([0, 1, 0], "zigzag"),      # Index 4 - longest
    ]

    env = SMDPfier(
        gym.make("CartPole-v1"),
        options_provider=options,
        duration_fn=ConstantOptionDuration(3),
        action_interface="index",
        max_options=5,
        availability_fn=adaptive_availability
    )

### Dynamic Options with Masking

When using dynamic option generators, the `availability_fn` is automatically passed to restrict generated options:

    from smdpfier.defaults.options import RandomStaticLen

    def state_aware_generator(obs, info):
        """Generate options based on state, using availability info."""
        # Get action mask from info (passed automatically)
        action_mask = info.get("action_mask")

        if action_mask is not None:
            # Generate options only with available actions
            available_actions = [i for i, avail in enumerate(action_mask) if avail]
        else:
            # No masking - use all actions
            available_actions = list(range(info["action_space"].n))

        # Generate random options with available actions only
        options = []
        for i in range(5):
            if available_actions:
                actions = random.choices(available_actions, k=3)
                options.append(Option(actions, f"dynamic_{i}"))

        return options

    def base_availability(obs):
        """Base availability function."""
        if obs[0] > 0:  # Cart right
            return [0]  # Only left action
        else:
            return [0, 1]  # Both actions

    env = SMDPfier(
        gym.make("CartPole-v1"),
        options_provider=state_aware_generator,
        duration_fn=ConstantOptionDuration(2),
        action_interface="index",
        max_options=5,
        availability_fn=base_availability  # Passed to generator automatically
    )

## Precheck Validation

Precheck validation attempts to validate options before execution by testing their actions in the current environment state.

### Enabling Precheck

    env = SMDPfier(
        base_env,
        options_provider=options,
        duration_fn=ConstantOptionDuration(5),
        action_interface="index",
        precheck=True  # Enable precheck validation
    )

### How Precheck Works

1. Before executing an option, SMDPfier saves the environment state
2. Tests each action in the option sequence
3. Restores the environment state after testing
4. Raises SMDPOptionValidationError if any action fails
5. Proceeds with execution if all actions are valid

### Precheck Example

    from smdpfier.errors import SMDPOptionValidationError

    # Option that might be invalid in some states
    risky_option = Option([0, 1, 0, 1, 0], "risky-sequence")

    try:
        obs, reward, term, trunc, info = env.step(risky_option)
    except SMDPOptionValidationError as e:
        print(f"Option '{e.option_name}' failed precheck!")
        print(f"Failed at step {e.failing_step_index}")
        print(f"Action {e.action_repr} is invalid")
        print(f"Environment state: {e.short_obs_summary}")

        # Handle the error (e.g., try a different option)
        fallback_option = Option([0], "safe-fallback")
        obs, reward, term, trunc, info = env.step(fallback_option)

### Precheck Limitations

⚠️ Important Limitations:

1. Environment must support state save/restore (not all environments do)
2. Performance overhead from testing each option
3. May not catch all edge cases (e.g., stochastic environments)
4. False positives possible in complex environments

Recommendation: Use precheck during development and debugging, consider disabling in production for performance.

### Precheck vs Masking

Approach | When to Use | Pros | Cons
--- | --- | --- | ---
Action Masking | Known invalid patterns | Fast, reliable | Requires domain knowledge
Precheck | Unknown failure modes | Automatic detection | Slower, may have false positives
Both | Maximum safety | Comprehensive validation | Higher complexity

## Best Practices

### Masking Strategy

    def robust_availability(obs):
        """Comprehensive availability function."""
        available = []

        # Conservative base set (always safe)
        available.extend([0, 1])  # Basic actions

        # Add options based on state confidence
        confidence = compute_state_confidence(obs)

        if confidence > 0.8:
            available.extend([2, 3])  # Medium complexity options

        if confidence > 0.9:
            available.extend([4, 5, 6])  # High complexity options

        return available

### Error-Resilient Option Design

    # Design options to minimize failure probability
    safe_options = [
        Option([0], "single-left"),       # Minimal option - rarely fails
        Option([1], "single-right"),      # Minimal option - rarely fails
        Option([0, 1], "balanced"),       # Balanced - self-correcting
    ]

    # Avoid overly long or extreme options
    risky_options = [
        Option([0]*10, "extreme-left"),   # Long sequence - high failure risk
        Option([1]*10, "extreme-right"),  # Long sequence - high failure risk
    ]

### Integration with RL Algorithms

    class MaskedDQNAgent:
        def select_action(self, obs, info):
            q_values = self.q_network(obs)

            # Apply action mask if available
            action_mask = info.get("smdp", {}).get("action_mask")
            if action_mask is not None:
                masked_q_values = q_values.copy()
                masked_q_values[np.array(action_mask) == 0] = -float('inf')
                return np.argmax(masked_q_values)
            else:
                return np.argmax(q_values)

        def train(self, env):
            obs, info = env.reset()

            while True:
                action = self.select_action(obs, info)

                try:
                    obs, reward, term, trunc, info = env.step(action)
                    # ... training logic ...
                except SMDPOptionValidationError:
                    # Handle validation failure
                    continue

                if term or trunc:
                    break

## Debugging Masking Issues

### Inspect Masking Behavior

    def debug_masking(env, num_steps=10):
        """Debug action masking behavior."""
        obs, info = env.reset()

        for step in range(num_steps):
            mask = info.get("smdp", {}).get("action_mask", [])
            available_actions = [i for i, avail in enumerate(mask) if avail == 1]

            print(f"Step {step}:")
            print(f"  Observation: {obs[:3]}...")  # First 3 elements
            print(f"  Action mask: {mask}")
            print(f"  Available actions: {available_actions}")

            if available_actions:
                action = random.choice(available_actions)
                obs, reward, term, trunc, info = env.step(action)

                if term or trunc:
                    break
            else:
                print("  No actions available!")
                break

    debug_masking(env)

### Common Masking Issues

Issue: All actions masked

    def overly_restrictive_availability(obs):
        if obs[0] > 2.0:  # Impossible condition
            return []  # No actions available!
        return [0, 1]

    # Fix: Ensure at least one action is always available
    def better_availability(obs):
        available = [0]  # Always allow basic action
        if obs[0] < 2.0:
            available.append(1)
        return available

Issue: Inconsistent masking

    def inconsistent_availability(obs):
        # Problem: Random masking
        return random.choices([0, 1, 2], k=random.randint(1, 3))

    # Fix: Deterministic masking based on state
    def consistent_availability(obs):
        position = obs[0]
        if position > 0.5:
            return [0, 2]  # Deterministic based on position
        else:
            return [1, 2]

## Summary

Feature | Purpose | Best For
--- | --- | ---
Action Masking | Restrict invalid options | State-dependent validity
Precheck | Test options before execution | Unknown failure modes
Combined | Maximum robustness | Critical applications

Key Takeaways: - Use masking for known patterns of invalid actions - Use precheck for unknown failure modes during development - Design options to minimize failure probability - Always ensure at least one action remains available

Next: Error Handling | See Also: API Reference