docstr: Returns a new input_dict after stepping through the environment once.
  
  Args:
      input_dict: the input dict passed into compute_single_action.
      action: the (unbatched) action taken this step.
      reward: the (unbatched) reward from env.step
      next_obs: the (unbatached) next observation from env.step
      extra: the extra action out from compute_single_action.
          In this case contains current returns to go *before* the current
          reward is subtracted from target_return.
  
  Returns:
      A new input_dict to be passed into compute_single_action.
      The input_dict for inference: {
          OBS: [max_seq_len, obs_dim] array,
          ACTIONS: [max_seq_len - 1, act_dim] array,
          RETURNS_TO_GO: [max_seq_len - 1] array,
          REWARDS: scalar,
          TIMESTEPS: [max_seq_len - 1] array,
      }
      Note the sequence lengths are different, and is specified as per
      view_requirements. Explanations in action_distribution_fn method.
docstr_tok: ["Returns","a","new","input_dict","after","stepping","through","the","environment","once",".","Args",":","input_dict",":","the","input","dict","passed","into","compute_single_action",".","action",":","the","(","unbatched",")","action","taken","this","step",".","reward",":","the","(","unbatched",")","reward","from","env.step","next_obs",":","the","(","unbatached",")","next","observation","from","env.step","extra",":","the","extra","action","out","from","compute_single_action",".","In","this","case","contains","current","returns","to","go","*","before","*","the","current","reward","is","subtracted","from","target_return",".","Returns",":","A","new","input_dict","to","be","passed","into","compute_single_action",".","The","input_dict","for","inference",":","{","OBS",":","[","max_seq_len",",","obs_dim","]","array",",","ACTIONS",":","[","max_seq_len","-","1",",","act_dim","]","array",",","RETURNS_TO_GO",":","[","max_seq_len","-","1","]","array",",","REWARDS",":","scalar",",","TIMESTEPS",":","[","max_seq_len","-","1","]","array",",","}","Note","the","sequence","lengths","are","different",",","and","is","specified","as","per","view_requirements",".","Explanations","in","action_distribution_fn","method","."]
code: def get_next_input_dict(
      self,
      input_dict: SampleBatch,
      action: TensorStructType,
      reward: TensorStructType,
      next_obs: TensorStructType,
      extra: Dict[str, TensorType],
  ) -> SampleBatch:
      """Returns a new input_dict after stepping through the environment once.
  
      Args:
          input_dict: the input dict passed into compute_single_action.
          action: the (unbatched) action taken this step.
          reward: the (unbatched) reward from env.step
          next_obs: the (unbatached) next observation from env.step
          extra: the extra action out from compute_single_action.
              In this case contains current returns to go *before* the current
              reward is subtracted from target_return.
  
      Returns:
          A new input_dict to be passed into compute_single_action.
          The input_dict for inference: {
              OBS: [max_seq_len, obs_dim] array,
              ACTIONS: [max_seq_len - 1, act_dim] array,
              RETURNS_TO_GO: [max_seq_len - 1] array,
              REWARDS: scalar,
              TIMESTEPS: [max_seq_len - 1] array,
          }
          Note the sequence lengths are different, and is specified as per
          view_requirements. Explanations in action_distribution_fn method.
      """
      # creates a copy of input_dict with only numpy arrays
      input_dict = tree.map_structure(convert_to_numpy, input_dict)
      # convert everything else to numpy as well
      action, reward, next_obs, extra = convert_to_numpy(
          (action, reward, next_obs, extra)
      )
  
      # check dimensions
      assert input_dict[SampleBatch.OBS].shape == (
          self.max_seq_len,
          *self.observation_space.shape,
      )
      assert input_dict[SampleBatch.ACTIONS].shape == (
          self.max_seq_len - 1,
          *self.action_space.shape,
      )
      assert input_dict[SampleBatch.RETURNS_TO_GO].shape == (
          self.max_seq_len - 1,
      )
      assert input_dict[SampleBatch.T].shape == (self.max_seq_len - 1,)
  
      # Shift observations
      input_dict[SampleBatch.OBS] = np.concatenate(
          [
              input_dict[SampleBatch.OBS][1:],
              next_obs[None],
          ],
          axis=0,
      )
  
      # Shift actions
      input_dict[SampleBatch.ACTIONS] = np.concatenate(
          [
              input_dict[SampleBatch.ACTIONS][1:],
              action[None],
          ],
          axis=0,
      )
  
      # Reward is not a sequence, it's only used to calculate next rtg.
      input_dict[SampleBatch.REWARDS] = np.asarray(reward)
  
      # See methods action_distribution_fn and extra_action_out for an explanation
      # of why this is done.
      input_dict[SampleBatch.RETURNS_TO_GO] = np.concatenate(
          [
              input_dict[SampleBatch.RETURNS_TO_GO][1:],
              np.asarray(extra[SampleBatch.RETURNS_TO_GO])[None],
          ],
          axis=0,
      )
  
      # Shift and increment timesteps
      input_dict[SampleBatch.T] = np.concatenate(
          [
              input_dict[SampleBatch.T][1:],
              input_dict[SampleBatch.T][-1:] + 1,
          ],
          axis=0,
      )
  
      return input_dict
  
imports: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Box, Discrete
  from rllib_dt.dt.dt_torch_model import DTTorchModel
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.mingpt import configure_gpt_optimizer
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDeterministic, TorchDistributionWrapper
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI, override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.threading import with_lock
  from ray.rllib.utils.torch_utils import apply_grad_clipping
  from ray.rllib.utils.typing import TensorShape, TensorStructType, TensorType, TrainerConfigDict
label_desc: Use image node_id 6 for calling the DTTorchPolicy obj's underlying member method code with example usage: obj.get_next_input_dict(input_dict, action, reward, next_obs, extra) and returns: input_dict