docstr: === DT configs
  
  Args:
      replay_buffer_config: Replay buffer config.
          Examples:
          {
          "_enable_replay_buffer_api": True,
          "type": "MultiAgentReplayBuffer",
          "capacity": 50000,
          "replay_sequence_length": 1,
          }
          - OR -
          {
          "_enable_replay_buffer_api": True,
          "type": "MultiAgentPrioritizedReplayBuffer",
          "capacity": 50000,
          "prioritized_replay_alpha": 0.6,
          "prioritized_replay_beta": 0.4,
          "prioritized_replay_eps": 1e-6,
          "replay_sequence_length": 1,
          }
          - Where -
          prioritized_replay_alpha: Alpha parameter controls the degree of
          prioritization in the buffer. In other words, when a buffer sample has
          a higher temporal-difference error, with how much more probability
          should it drawn to use to update the parametrized Q-network. 0.0
          corresponds to uniform probability. Setting much above 1.0 may quickly
          result as the sampling distribution could become heavily “pointy” with
          low entropy.
          prioritized_replay_beta: Beta parameter controls the degree of
          importance sampling which suppresses the influence of gradient updates
          from samples that have higher probability of being sampled via alpha
          parameter and the temporal-difference error.
          prioritized_replay_eps: Epsilon parameter sets the baseline probability
          for sampling so that when the temporal-difference error of a sample is
          zero, there is still a chance of drawing the sample.
      embed_dim: Dimension of the embeddings in the GPT model.
      num_layers: Number of attention layers in the GPT model.
      num_heads: Number of attention heads in the GPT model. Must divide
          embed_dim evenly.
      embed_pdrop: Dropout probability of the embedding layer of the GPT model.
      resid_pdrop: Dropout probability of the residual layer of the GPT model.
      attn_pdrop: Dropout probability of the attention layer of the GPT model.
      grad_clip: If specified, clip the global norm of gradients by this amount.
      lr_schedule: Learning rate schedule. In the format of
          [[timestep, lr-value], [timestep, lr-value], ...]
          Intermediary timesteps will be assigned to interpolated learning rate
          values. A schedule should normally start from timestep 0.
      loss_coef_actions: Coefficients on the loss for the actions output.
          Default to 1.
      loss_coef_obs: Coefficients on the loss for the obs output. Default to 0.
          Set to a value greater than 0 to regress on the obs output.
      loss_coef_returns_to_go: Coefficients on the loss for the returns_to_go
          output. Default to 0. Set to a value greater than 0 to regress on the
          returns_to_go output.
      horizon: The episode horizon used. This value can be derived from your
          environment via `[your_env]._max_episode_steps`.
      **kwargs: Forward compatibility kwargs
  
  Returns:
      This updated DTConfig object.
docstr_tok: ["===","DT","configs","Args",":","replay_buffer_config",":","Replay","buffer","config",".","Examples",":","{","``","_enable_replay_buffer_api","''",":","True",",","``","type","''",":","``","MultiAgentReplayBuffer","''",",","``","capacity","''",":","50000",",","``","replay_sequence_length","''",":","1",",","}","-","OR","-","{","``","_enable_replay_buffer_api","''",":","True",",","``","type","''",":","``","MultiAgentPrioritizedReplayBuffer","''",",","``","capacity","''",":","50000",",","``","prioritized_replay_alpha","''",":","0.6",",","``","prioritized_replay_beta","''",":","0.4",",","``","prioritized_replay_eps","''",":","1e-6",",","``","replay_sequence_length","''",":","1",",","}","-","Where","-","prioritized_replay_alpha",":","Alpha","parameter","controls","the","degree","of","prioritization","in","the","buffer",".","In","other","words",",","when","a","buffer","sample","has","a","higher","temporal-difference","error",",","with","how","much","more","probability","should","it","drawn","to","use","to","update","the","parametrized","Q-network",".","0.0","corresponds","to","uniform","probability",".","Setting","much","above","1.0","may","quickly","result","as","the","sampling","distribution","could","become","heavily","\u201c","pointy","\u201d","with","low","entropy",".","prioritized_replay_beta",":","Beta","parameter","controls","the","degree","of","importance","sampling","which","suppresses","the","influence","of","gradient","updates","from","samples","that","have","higher","probability","of","being","sampled","via","alpha","parameter","and","the","temporal-difference","error",".","prioritized_replay_eps",":","Epsilon","parameter","sets","the","baseline","probability","for","sampling","so","that","when","the","temporal-difference","error","of","a","sample","is","zero",",","there","is","still","a","chance","of","drawing","the","sample",".","embed_dim",":","Dimension","of","the","embeddings","in","the","GPT","model",".","num_layers",":","Number","of","attention","layers","in","the","GPT","model",".","num_heads",":","Number","of","attention","heads","in","the","GPT","model",".","Must","divide","embed_dim","evenly",".","embed_pdrop",":","Dropout","probability","of","the","embedding","layer","of","the","GPT","model",".","resid_pdrop",":","Dropout","probability","of","the","residual","layer","of","the","GPT","model",".","attn_pdrop",":","Dropout","probability","of","the","attention","layer","of","the","GPT","model",".","grad_clip",":","If","specified",",","clip","the","global","norm","of","gradients","by","this","amount",".","lr_schedule",":","Learning","rate","schedule",".","In","the","format","of","[","[","timestep",",","lr-value","]",",","[","timestep",",","lr-value","]",",","...","]","Intermediary","timesteps","will","be","assigned","to","interpolated","learning","rate","values",".","A","schedule","should","normally","start","from","timestep","0.","loss_coef_actions",":","Coefficients","on","the","loss","for","the","actions","output",".","Default","to","1.","loss_coef_obs",":","Coefficients","on","the","loss","for","the","obs","output",".","Default","to","0",".","Set","to","a","value","greater","than","0","to","regress","on","the","obs","output",".","loss_coef_returns_to_go",":","Coefficients","on","the","loss","for","the","returns_to_go","output",".","Default","to","0",".","Set","to","a","value","greater","than","0","to","regress","on","the","returns_to_go","output",".","horizon",":","The","episode","horizon","used",".","This","value","can","be","derived","from","your","environment","via","`","[","your_env","]","._max_episode_steps","`",".","*","*","kwargs",":","Forward","compatibility","kwargs","Returns",":","This","updated","DTConfig","object","."]
code: def training(
      self,
      *,
      replay_buffer_config: Optional[Dict[str, Any]] = NotProvided,
      embed_dim: Optional[int] = NotProvided,
      num_layers: Optional[int] = NotProvided,
      num_heads: Optional[int] = NotProvided,
      embed_pdrop: Optional[float] = NotProvided,
      resid_pdrop: Optional[float] = NotProvided,
      attn_pdrop: Optional[float] = NotProvided,
      grad_clip: Optional[float] = NotProvided,
      loss_coef_actions: Optional[float] = NotProvided,
      loss_coef_obs: Optional[float] = NotProvided,
      loss_coef_returns_to_go: Optional[float] = NotProvided,
      lr_schedule: Optional[
          List[List[Union[int, float]]]
      ] = NotProvided,
      horizon: Optional[int] = NotProvided,
      **kwargs,
  ) -> "DTConfig":
      """
      === DT configs
  
      Args:
          replay_buffer_config: Replay buffer config.
              Examples:
              {
              "_enable_replay_buffer_api": True,
              "type": "MultiAgentReplayBuffer",
              "capacity": 50000,
              "replay_sequence_length": 1,
              }
              - OR -
              {
              "_enable_replay_buffer_api": True,
              "type": "MultiAgentPrioritizedReplayBuffer",
              "capacity": 50000,
              "prioritized_replay_alpha": 0.6,
              "prioritized_replay_beta": 0.4,
              "prioritized_replay_eps": 1e-6,
              "replay_sequence_length": 1,
              }
              - Where -
              prioritized_replay_alpha: Alpha parameter controls the degree of
              prioritization in the buffer. In other words, when a buffer sample has
              a higher temporal-difference error, with how much more probability
              should it drawn to use to update the parametrized Q-network. 0.0
              corresponds to uniform probability. Setting much above 1.0 may quickly
              result as the sampling distribution could become heavily “pointy” with
              low entropy.
              prioritized_replay_beta: Beta parameter controls the degree of
              importance sampling which suppresses the influence of gradient updates
              from samples that have higher probability of being sampled via alpha
              parameter and the temporal-difference error.
              prioritized_replay_eps: Epsilon parameter sets the baseline probability
              for sampling so that when the temporal-difference error of a sample is
              zero, there is still a chance of drawing the sample.
          embed_dim: Dimension of the embeddings in the GPT model.
          num_layers: Number of attention layers in the GPT model.
          num_heads: Number of attention heads in the GPT model. Must divide
              embed_dim evenly.
          embed_pdrop: Dropout probability of the embedding layer of the GPT model.
          resid_pdrop: Dropout probability of the residual layer of the GPT model.
          attn_pdrop: Dropout probability of the attention layer of the GPT model.
          grad_clip: If specified, clip the global norm of gradients by this amount.
          lr_schedule: Learning rate schedule. In the format of
              [[timestep, lr-value], [timestep, lr-value], ...]
              Intermediary timesteps will be assigned to interpolated learning rate
              values. A schedule should normally start from timestep 0.
          loss_coef_actions: Coefficients on the loss for the actions output.
              Default to 1.
          loss_coef_obs: Coefficients on the loss for the obs output. Default to 0.
              Set to a value greater than 0 to regress on the obs output.
          loss_coef_returns_to_go: Coefficients on the loss for the returns_to_go
              output. Default to 0. Set to a value greater than 0 to regress on the
              returns_to_go output.
          horizon: The episode horizon used. This value can be derived from your
              environment via `[your_env]._max_episode_steps`.
          **kwargs: Forward compatibility kwargs
  
      Returns:
          This updated DTConfig object.
      """
      super().training(**kwargs)
      if replay_buffer_config is not NotProvided:
          # Override entire `replay_buffer_config` if `type` key changes.
          # Update, if `type` key remains the same or is not specified.
          new_replay_buffer_config = deep_update(
              {"replay_buffer_config": self.replay_buffer_config},
              {"replay_buffer_config": replay_buffer_config},
              False,
              ["replay_buffer_config"],
              ["replay_buffer_config"],
          )
          self.replay_buffer_config = new_replay_buffer_config[
              "replay_buffer_config"
          ]
      if embed_dim is not NotProvided:
          self.embed_dim = embed_dim
      if num_layers is not NotProvided:
          self.num_layers = num_layers
      if num_heads is not NotProvided:
          self.num_heads = num_heads
      if embed_pdrop is not NotProvided:
          self.embed_pdrop = embed_pdrop
      if resid_pdrop is not NotProvided:
          self.resid_pdrop = resid_pdrop
      if attn_pdrop is not NotProvided:
          self.attn_pdrop = attn_pdrop
      if grad_clip is not NotProvided:
          self.grad_clip = grad_clip
      if lr_schedule is not NotProvided:
          self.lr_schedule = lr_schedule
      if loss_coef_actions is not NotProvided:
          self.loss_coef_actions = loss_coef_actions
      if loss_coef_obs is not NotProvided:
          self.loss_coef_obs = loss_coef_obs
      if loss_coef_returns_to_go is not NotProvided:
          self.loss_coef_returns_to_go = loss_coef_returns_to_go
      if horizon is not NotProvided:
          self.horizon = horizon
  
      return self
  
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 2 for calling the DTConfig obj's underlying member method code with example usage: obj.training() and returns: self