docstr: Test loss function.
docstr_tok: ["Test","loss","function","."]
code: def test_loss(self):
      """Test loss function."""
      config = _default_config()
      config["embed_pdrop"] = 0
      config["resid_pdrop"] = 0
      config["attn_pdrop"] = 0
  
      observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,))
      action_spaces = [
          gym.spaces.Box(-1.0, 1.0, shape=(1,)),
          gym.spaces.Discrete(4),
      ]
  
      for action_space in action_spaces:
          # Create policy
          policy = DTTorchPolicy(
              observation_space, action_space, config
          )
  
          # Run loss functions on batches with different items in the mask to make
          # sure the masks are working and making the loss the same.
          batch1 = SampleBatch(
              {
                  SampleBatch.OBS: np.array(
                      [
                          [
                              [0.0, 0.0, 0.0],
                              [0.0, 0.0, 0.0],
                              [0.0, 1.0, 2.0],
                              [3.0, 4.0, 5.0],
                          ]
                      ],
                      dtype=np.float32,
                  ),
                  SampleBatch.ACTIONS: (
                      np.array(
                          [[[0.0], [0.0], [1.0], [0.5]]],
                          dtype=np.float32,
                      )
                      if isinstance(action_space, gym.spaces.Box)
                      else np.array([[0, 0, 1, 3]], dtype=np.int64)
                  ),
                  SampleBatch.RETURNS_TO_GO: np.array(
                      [[[0.0], [0.0], [100.0], [90.0], [80.0]]],
                      dtype=np.float32,
                  ),
                  SampleBatch.T: np.array(
                      [[0, 0, 0, 1]], dtype=np.int32
                  ),
                  SampleBatch.ATTENTION_MASKS: np.array(
                      [[0.0, 0.0, 1.0, 1.0]], dtype=np.float32
                  ),
              }
          )
  
          batch2 = SampleBatch(
              {
                  SampleBatch.OBS: np.array(
                      [
                          [
                              [1.0, 1.0, -1.0],
                              [1.0, 10.0, 12.0],
                              [0.0, 1.0, 2.0],
                              [3.0, 4.0, 5.0],
                          ]
                      ],
                      dtype=np.float32,
                  ),
                  SampleBatch.ACTIONS: (
                      np.array(
                          [[[1.0], [-0.5], [1.0], [0.5]]],
                          dtype=np.float32,
                      )
                      if isinstance(action_space, gym.spaces.Box)
                      else np.array([[2, 1, 1, 3]], dtype=np.int64)
                  ),
                  SampleBatch.RETURNS_TO_GO: np.array(
                      [[[200.0], [-10.0], [100.0], [90.0], [80.0]]],
                      dtype=np.float32,
                  ),
                  SampleBatch.T: np.array(
                      [[9, 3, 0, 1]], dtype=np.int32
                  ),
                  SampleBatch.ATTENTION_MASKS: np.array(
                      [[0.0, 0.0, 1.0, 1.0]], dtype=np.float32
                  ),
              }
          )
  
          loss1 = policy.loss(policy.model, policy.dist_class, batch1)
          loss2 = policy.loss(policy.model, policy.dist_class, batch2)
  
          loss1 = loss1.detach().cpu().item()
          loss2 = loss2.detach().cpu().item()
  
          assert np.isclose(
              loss1, loss2
          ), "Masks are not working for losses."
  
          # Run loss on a widely different batch and make sure the loss is different.
          batch3 = SampleBatch(
              {
                  SampleBatch.OBS: np.array(
                      [
                          [
                              [1.0, 1.0, -20.0],
                              [0.1, 10.0, 12.0],
                              [1.4, 12.0, -9.0],
                              [6.0, 40.0, -2.0],
                          ]
                      ],
                      dtype=np.float32,
                  ),
                  SampleBatch.ACTIONS: (
                      np.array(
                          [[[2.0], [-1.5], [0.2], [0.1]]],
                          dtype=np.float32,
                      )
                      if isinstance(action_space, gym.spaces.Box)
                      else np.array([[1, 3, 0, 2]], dtype=np.int64)
                  ),
                  SampleBatch.RETURNS_TO_GO: np.array(
                      [[[90.0], [80.0], [70.0], [60.0], [50.0]]],
                      dtype=np.float32,
                  ),
                  SampleBatch.T: np.array(
                      [[3, 4, 5, 6]], dtype=np.int32
                  ),
                  SampleBatch.ATTENTION_MASKS: np.array(
                      [[1.0, 1.0, 1.0, 1.0]], dtype=np.float32
                  ),
              }
          )
  
          loss3 = policy.loss(policy.model, policy.dist_class, batch3)
          loss3 = loss3.detach().cpu().item()
  
          assert not np.isclose(
              loss1, loss3
          ), "Widely different inputs are giving the same loss value."
  
file_path: ray/rllib_contrib/dt/tests/test_dt_policy.py
imports: import unittest
  from typing import Dict
  import gymnasium
  import numpy
  from rllib_dt.dt.dt_torch_policy import DTTorchPolicy
  import ray
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
label_desc: Use image node_id 6 for calling the TestDTPolicy obj's underlying member method code with example usage: obj.test_loss() without return types