class_docstr: Implements Decision Transformer: https://arxiv.org/abs/2106.01345.
class_docstr_tok: ["Implements","Decision","Transformer",":","https",":","\/\/arxiv.org\/abs\/2106.01345","."]
code: def training_step(self) -> ResultDict:
      with self._timers[SAMPLE_TIMER]:
          # TODO: Add ability to do obs_filter for offline sampling.
          train_batch = synchronous_parallel_sample(
              worker_set=self.workers
          )
  
      train_batch = train_batch.as_multi_agent()
      self._counters[
          NUM_AGENT_STEPS_SAMPLED
      ] += train_batch.agent_steps()
      self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps()
  
      # Because each sample is a segment of max_seq_len transitions, doing
      # the division makes it so the total number of transitions per train
      # step is consistent.
      num_steps = train_batch.env_steps()
      batch_size = int(
          math.ceil(num_steps / self.config.model["max_seq_len"])
      )
  
      # Add the batch of episodes to the segmentation buffer.
      self.local_replay_buffer.add(train_batch)
      # Sample a batch of segments.
      train_batch = self.local_replay_buffer.sample(batch_size)
  
      # Postprocess batch before we learn on it.
      post_fn = self.config.get("before_learn_on_batch") or (
          lambda b, *a: b
      )
      train_batch = post_fn(train_batch, self.workers, self.config)
  
      # Learn on training batch.
      # Use simple optimizer (only for multi-agent or tf-eager; all other
      # cases should use the multi-GPU optimizer, even if only using 1 GPU)
      if self.config.get("simple_optimizer", False):
          train_results = train_one_step(self, train_batch)
      else:
          train_results = multi_gpu_train_one_step(self, train_batch)
  
      # Update learning rate scheduler.
      global_vars = {
          # Note: this counts the number of segments trained, not timesteps.
          # i.e. NUM_AGENT_STEPS_TRAINED: B, NUM_AGENT_STEPS_SAMPLED: B*T
          "timestep": self._counters[NUM_AGENT_STEPS_TRAINED],
      }
      self.workers.local_worker().set_global_vars(global_vars)
  
      return train_results
  
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 3 for calling the DT obj's underlying member method code with example usage: obj.training_step() and returns: train_results