docstr: Test inference input_dict methods
  
  This is a minimal version the test in test_dt.py.
  The shapes of the input_dict might be confusing but it makes sense in
  context of what the function is supposed to do.
  Check action_distribution_fn for an explanation.
docstr_tok: ["Test","inference","input_dict","methods","This","is","a","minimal","version","the","test","in","test_dt.py",".","The","shapes","of","the","input_dict","might","be","confusing","but","it","makes","sense","in","context","of","what","the","function","is","supposed","to","do",".","Check","action_distribution_fn","for","an","explanation","."]
code: def test_torch_input_dict(self):
      """Test inference input_dict methods
  
      This is a minimal version the test in test_dt.py.
      The shapes of the input_dict might be confusing but it makes sense in
      context of what the function is supposed to do.
      Check action_distribution_fn for an explanation.
      """
      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
          )
  
          # initial obs and input_dict
          obs = np.array([0.0, 1.0, 2.0])
          input_dict = policy.get_initial_input_dict(obs)
  
          # Check input_dict matches what it should be
          target_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.zeros((), dtype=np.float32),
                  SampleBatch.T: np.array([-1, -1, -1], dtype=np.int32),
              }
          )
          _assert_input_dict_equals(input_dict, target_input_dict)
  
          # Get next input_dict
          input_dict = policy.get_next_input_dict(
              input_dict,
              action=(
                  np.asarray([1.0], dtype=np.float32)
                  if isinstance(action_space, gym.spaces.Box)
                  else np.asarray(1, dtype=np.int32)
              ),
              reward=1.0,
              next_obs=np.array([3.0, 4.0, 5.0]),
              extra={
                  SampleBatch.RETURNS_TO_GO: config["target_return"],
              },
          )
  
          # Check input_dict matches what it should be
          target_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.asarray(
                      1.0, dtype=np.float32
                  ),
                  SampleBatch.T: np.array([-1, -1, 0], dtype=np.int32),
              }
          )
          _assert_input_dict_equals(input_dict, target_input_dict)
  
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 4 for calling the TestDTPolicy obj's underlying member method code with example usage: obj.test_torch_input_dict() without return types