class_docstr: Implements Decision Transformer: https://arxiv.org/abs/2106.01345.
class_docstr_tok: ["Implements","Decision","Transformer",":","https",":","\/\/arxiv.org\/abs\/2106.01345","."]
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.
          For DT this case contains current returns to go *before* the current
          reward is subtracted from target_return.
      policy_id: Policy to query (only applies to multi-agent).
          Default: "default_policy".
  
  Returns:
      A new input_dict to be passed into compute_single_action.
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",".","For","DT","this","case","contains","current","returns","to","go","*","before","*","the","current","reward","is","subtracted","from","target_return",".","policy_id",":","Policy","to","query","(","only","applies","to","multi-agent",")",".","Default",":","``","default_policy","''",".","Returns",":","A","new","input_dict","to","be","passed","into","compute_single_action","."]
code: def get_next_input_dict(
      self,
      input_dict: SampleBatch,
      action: TensorStructType,
      reward: TensorStructType,
      next_obs: TensorStructType,
      extra: Dict[str, TensorType],
      policy_id: PolicyID = DEFAULT_POLICY_ID,
  ) -> 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.
              For DT this case contains current returns to go *before* the current
              reward is subtracted from target_return.
          policy_id: Policy to query (only applies to multi-agent).
              Default: "default_policy".
  
      Returns:
          A new input_dict to be passed into compute_single_action.
      """
      policy = self.get_policy(policy_id)
      return policy.get_next_input_dict(
          input_dict, action, reward, next_obs, extra
      )
  
imports: import logging
  import math
  from typing import Any, Dict, List, Optional, Tuple, Type, Union
  from rllib_dt.dt.segmentation_buffer import MultiAgentSegmentationBuffer
  from ray.rllib import SampleBatch
  from ray.rllib.algorithms.algorithm import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.execution import synchronous_parallel_sample
  from ray.rllib.execution.train_ops import multi_gpu_train_one_step, train_one_step
  from ray.rllib.policy import Policy
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils import deep_update
  from ray.rllib.utils.annotations import PublicAPI, override
  from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_AGENT_STEPS_TRAINED, NUM_ENV_STEPS_SAMPLED, SAMPLE_TIMER
  from ray.rllib.utils.typing import PolicyID, ResultDict, TensorStructType, TensorType
label_desc: Use image node_id 6 for calling the DT obj's underlying member method code with example usage: obj.get_next_input_dict(input_dict, action, reward, next_obs, extra, policy_id) and returns: policy