docstr: Test that the transformer model' causal masking works.
docstr_tok: ["Test","that","the","transformer","model","'","causal","masking","works","."]
code: def test_causal_masking(self):
      """Test that the transformer model' causal masking works."""
      model_config = {
          "embed_dim": 16,
          "num_layers": 2,
          "max_seq_len": 4,
          "max_ep_len": 10,
          "num_heads": 2,
          "embed_pdrop": 0,
          "resid_pdrop": 0,
          "attn_pdrop": 0,
          "use_obs_output": True,
          "use_return_output": True,
      }
  
      observation_space = gym.spaces.Box(-1.0, 1.0, shape=(4,))
      action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,))
      B = 2
      T = model_config["max_seq_len"]
  
      # Generate input dict.
      input_dict = _generate_input_dict(
          B, T, observation_space, action_space
      )
  
      # make model and forward with attention
      model = DTTorchModel(
          observation_space,
          action_space,
          4,
          model_config,
          "model",
      )
      model_out, _ = model(input_dict)
      preds = model.get_prediction(
          model_out, input_dict, return_attentions=True
      )
      preds = convert_to_numpy(preds)
  
      # test properties of attentions
      attentions = preds["attentions"]
      self.assertEqual(
          len(attentions),
          model_config["num_layers"],
          "there should as many attention tensors as layers.",
      )
  
      # used to select the causal padded element of each attention tensor
      select_mask = np.triu(
          np.ones((3 * T, 3 * T), dtype=np.bool_), k=1
      )
      select_mask = np.tile(
          select_mask, (B, model_config["num_heads"], 1, 1)
      )
  
      for attention in attentions:
          # check shape
          self.assertEqual(
              attention.shape,
              (B, model_config["num_heads"], T * 3, T * 3),
          )
          # check the upper triangular masking
          assert np.allclose(
              attention[select_mask], 0.0
          ), "masked elements should be zero."
          # check that the non-masked elements have non 0 scores
          # Note: it is very unlikely that randomly initialized weights will make
          # one of the scores be 0, as these scores are probabilities.
          assert not np.any(
              np.isclose(attention[np.logical_not(select_mask)], 0.0)
          ), "non masked elements should be nonzero."
  
file_path: ray/rllib_contrib/dt/tests/test_dt_model.py
imports: import unittest
  import gymnasium
  import numpy
  from rllib_dt.dt.dt_torch_model import DTTorchModel
  import ray
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
label_desc: Use image node_id 5 for calling the TestDTModel obj's underlying member method code with example usage: obj.test_causal_masking() without return types