class_docstr: Model that handles simple discrete action masking.
  
  This assumes the outputs are logits for a single Categorical action dist.
  Getting this to work with a more complex output (e.g., if the action space
  is a tuple of several distributions) is also possible but left as an
  exercise to the reader.
class_docstr_tok: ["Model","that","handles","simple","discrete","action","masking",".","This","assumes","the","outputs","are","logits","for","a","single","Categorical","action","dist",".","Getting","this","to","work","with","a","more","complex","output","(","e.g.",",","if","the","action","space","is","a","tuple","of","several","distributions",")","is","also","possible","but","left","as","an","exercise","to","the","reader","."]
code: def forward(self, input_dict, state, seq_lens):
      # Extract the available actions tensor from the observation.
      action_mask = input_dict["obs"]["action_mask"]
  
      # Compute the unmasked logits.
      logits, _ = self.internal_model(
          {"obs": input_dict["obs"]["observations"]}
      )
  
      # If action masking is disabled, directly return unmasked logits
      if self.no_masking:
          return logits, state
  
      # Convert action_mask into a [0.0 || -inf]-type mask.
      inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)
      masked_logits = logits + inf_mask
  
      # Return masked logits.
      return masked_logits, state
  
file_path: ray/rllib/examples/models/action_mask_model.py
imports: from gymnasium.spaces import Dict
  from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.models.torch.fcnet import FullyConnectedNetwork
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.torch_utils import FLOAT_MIN
label_desc: Use image node_id 2 for calling the ActionMaskModel obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: masked_logits,  state, logits,  state