docstr: Test policy's action_distribution_fn and extra_action_out methods by
  calling compute_actions_from_input_dict which works those two methods
  in conjunction.
docstr_tok: ["Test","policy","'s","action_distribution_fn","and","extra_action_out","methods","by","calling","compute_actions_from_input_dict","which","works","those","two","methods","in","conjunction","."]
code: def test_torch_action(self):
      """Test policy's action_distribution_fn and extra_action_out methods by
      calling compute_actions_from_input_dict which works those two methods
      in conjunction.
      """
      config = _default_config()
  
      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
          )
  
          # input_dict for initial observation
          input_dict = SampleBatch(
              {
                  SampleBatch.OBS: np.array(
                      [
                          [
                              [0.0, 0.0, 0.0],
                              [0.0, 0.0, 0.0],
                              [0.0, 0.0, 0.0],
                              [0.0, 1.0, 2.0],
                          ]
                      ],
                      dtype=np.float32,
                  ),
                  SampleBatch.ACTIONS: (
                      np.array(
                          [[[0.0], [0.0], [0.0]]], dtype=np.float32
                      )
                      if isinstance(action_space, gym.spaces.Box)
                      else np.array([[0, 0, 0]], dtype=np.int32)
                  ),
                  SampleBatch.RETURNS_TO_GO: np.array(
                      [[0.0, 0.0, 0.0]], dtype=np.float32
                  ),
                  SampleBatch.REWARDS: np.array(
                      [0.0], dtype=np.float32
                  ),
                  SampleBatch.T: np.array(
                      [[-1, -1, -1]], dtype=np.int32
                  ),
              }
          )
  
          # Run compute_actions_from_input_dict
          actions, _, extras = policy.compute_actions_from_input_dict(
              input_dict,
              explore=False,
              timestep=None,
          )
  
          # Check actions
          assert actions.shape == (
              1,
              *action_space.shape,
          ), "actions has incorrect shape."
  
          # Check extras
          assert (
              SampleBatch.RETURNS_TO_GO in extras
          ), "extras should contain returns_to_go."
          assert extras[SampleBatch.RETURNS_TO_GO].shape == (
              1,
          ), "extras['returns_to_go'] has incorrect shape."
          assert np.isclose(
              extras[SampleBatch.RETURNS_TO_GO],
              np.asarray([config["target_return"]], dtype=np.float32),
          ), "extras['returns_to_go'] should contain target_return."
  
          # input_dict for non-initial observation
          input_dict = 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]]], dtype=np.float32
                      )
                      if isinstance(action_space, gym.spaces.Box)
                      else np.array([[0, 0, 1]], dtype=np.int32)
                  ),
                  SampleBatch.RETURNS_TO_GO: np.array(
                      [[0.0, 0.0, config["target_return"]]],
                      dtype=np.float32,
                  ),
                  SampleBatch.REWARDS: np.array(
                      [10.0], dtype=np.float32
                  ),
                  SampleBatch.T: np.array(
                      [[-1, -1, 0]], dtype=np.int32
                  ),
              }
          )
  
          # Run compute_actions_from_input_dict
          actions, _, extras = policy.compute_actions_from_input_dict(
              input_dict,
              explore=False,
              timestep=None,
          )
  
          # Check actions
          assert actions.shape == (
              1,
              *action_space.shape,
          ), "actions has incorrect shape."
  
          # Check extras
          assert (
              SampleBatch.RETURNS_TO_GO in extras
          ), "extras should contain returns_to_go."
          assert extras[SampleBatch.RETURNS_TO_GO].shape == (
              1,
          ), "extras['returns_to_go'] has incorrect shape."
          assert np.isclose(
              extras[SampleBatch.RETURNS_TO_GO],
              np.asarray(
                  [config["target_return"] - 10.0], dtype=np.float32
              ),
          ), "extras['returns_to_go'] should contain target_return."
  
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 5 for calling the TestDTPolicy obj's underlying member method code with example usage: obj.test_torch_action() without return types