class_docstr: PyTorch specific Policy class to use with RLlib.
class_docstr_tok: ["PyTorch","specific","Policy","class","to","use","with","RLlib","."]
docstr: Initializes a TorchPolicy instance.
  
  Args:
      observation_space: Observation space of the policy.
      action_space: Action space of the policy.
      config: The Policy's config dict.
      model: PyTorch policy module. Given observations as
          input, this module must return a list of outputs where the
          first item is action logits, and the rest can be any value.
      loss: Callable that returns one or more (a list of) scalar loss
          terms.
      action_distribution_class: Class for a torch action distribution.
      action_sampler_fn: A callable returning either a sampled action,
          its log-likelihood and updated state or a sampled action, its
          log-likelihood, updated state and action distribution inputs
          given Policy, ModelV2, input_dict, state batches (optional),
          explore, and timestep. Provide `action_sampler_fn` if you would
          like to have full control over the action computation step,
          including the model forward pass, possible sampling from a
          distribution, and exploration logic.
          Note: If `action_sampler_fn` is given, `action_distribution_fn`
          must be None. If both `action_sampler_fn` and
          `action_distribution_fn` are None, RLlib will simply pass
          inputs through `self.model` to get distribution inputs, create
          the distribution object, sample from it, and apply some
          exploration logic to the results.
          The callable takes as inputs: Policy, ModelV2, input_dict
          (SampleBatch), state_batches (optional), explore, and timestep.
      action_distribution_fn: A callable returning distribution inputs
          (parameters), a dist-class to generate an action distribution
          object from, and internal-state outputs (or an empty list if
          not applicable).
          Provide `action_distribution_fn` if you would like to only
          customize the model forward pass call. The resulting
          distribution parameters are then used by RLlib to create a
          distribution object, sample from it, and execute any
          exploration logic.
          Note: If `action_distribution_fn` is given, `action_sampler_fn`
          must be None. If both `action_sampler_fn` and
          `action_distribution_fn` are None, RLlib will simply pass
          inputs through `self.model` to get distribution inputs, create
          the distribution object, sample from it, and apply some
          exploration logic to the results.
          The callable takes as inputs: Policy, ModelV2, ModelInputDict,
          explore, timestep, is_training.
      max_seq_len: Max sequence length for LSTM training.
      get_batch_divisibility_req: Optional callable that returns the
          divisibility requirement for sample batches given the Policy.
docstr_tok: ["Initializes","a","TorchPolicy","instance",".","Args",":","observation_space",":","Observation","space","of","the","policy",".","action_space",":","Action","space","of","the","policy",".","config",":","The","Policy","'s","config","dict",".","model",":","PyTorch","policy","module",".","Given","observations","as","input",",","this","module","must","return","a","list","of","outputs","where","the","first","item","is","action","logits",",","and","the","rest","can","be","any","value",".","loss",":","Callable","that","returns","one","or","more","(","a","list","of",")","scalar","loss","terms",".","action_distribution_class",":","Class","for","a","torch","action","distribution",".","action_sampler_fn",":","A","callable","returning","either","a","sampled","action",",","its","log-likelihood","and","updated","state","or","a","sampled","action",",","its","log-likelihood",",","updated","state","and","action","distribution","inputs","given","Policy",",","ModelV2",",","input_dict",",","state","batches","(","optional",")",",","explore",",","and","timestep",".","Provide","`","action_sampler_fn","`","if","you","would","like","to","have","full","control","over","the","action","computation","step",",","including","the","model","forward","pass",",","possible","sampling","from","a","distribution",",","and","exploration","logic",".","Note",":","If","`","action_sampler_fn","`","is","given",",","`","action_distribution_fn","`","must","be","None",".","If","both","`","action_sampler_fn","`","and","`","action_distribution_fn","`","are","None",",","RLlib","will","simply","pass","inputs","through","`","self.model","`","to","get","distribution","inputs",",","create","the","distribution","object",",","sample","from","it",",","and","apply","some","exploration","logic","to","the","results",".","The","callable","takes","as","inputs",":","Policy",",","ModelV2",",","input_dict","(","SampleBatch",")",",","state_batches","(","optional",")",",","explore",",","and","timestep",".","action_distribution_fn",":","A","callable","returning","distribution","inputs","(","parameters",")",",","a","dist-class","to","generate","an","action","distribution","object","from",",","and","internal-state","outputs","(","or","an","empty","list","if","not","applicable",")",".","Provide","`","action_distribution_fn","`","if","you","would","like","to","only","customize","the","model","forward","pass","call",".","The","resulting","distribution","parameters","are","then","used","by","RLlib","to","create","a","distribution","object",",","sample","from","it",",","and","execute","any","exploration","logic",".","Note",":","If","`","action_distribution_fn","`","is","given",",","`","action_sampler_fn","`","must","be","None",".","If","both","`","action_sampler_fn","`","and","`","action_distribution_fn","`","are","None",",","RLlib","will","simply","pass","inputs","through","`","self.model","`","to","get","distribution","inputs",",","create","the","distribution","object",",","sample","from","it",",","and","apply","some","exploration","logic","to","the","results",".","The","callable","takes","as","inputs",":","Policy",",","ModelV2",",","ModelInputDict",",","explore",",","timestep",",","is_training",".","max_seq_len",":","Max","sequence","length","for","LSTM","training",".","get_batch_divisibility_req",":","Optional","callable","that","returns","the","divisibility","requirement","for","sample","batches","given","the","Policy","."]
code: def __init__(
      self,
      observation_space: gym.spaces.Space,
      action_space: gym.spaces.Space,
      config: AlgorithmConfigDict,
      *,
      model: Optional[TorchModelV2] = None,
      loss: Optional[
          Callable[
              [
                  Policy,
                  ModelV2,
                  Type[TorchDistributionWrapper],
                  SampleBatch,
              ],
              Union[TensorType, List[TensorType]],
          ]
      ] = None,
      action_distribution_class: Optional[
          Type[TorchDistributionWrapper]
      ] = None,
      action_sampler_fn: Optional[
          Callable[
              [TensorType, List[TensorType]],
              Union[
                  Tuple[TensorType, TensorType, List[TensorType]],
                  Tuple[
                      TensorType,
                      TensorType,
                      TensorType,
                      List[TensorType],
                  ],
              ],
          ]
      ] = None,
      action_distribution_fn: Optional[
          Callable[
              [Policy, ModelV2, TensorType, TensorType, TensorType],
              Tuple[
                  TensorType,
                  Type[TorchDistributionWrapper],
                  List[TensorType],
              ],
          ]
      ] = None,
      max_seq_len: int = 20,
      get_batch_divisibility_req: Optional[
          Callable[[Policy], int]
      ] = None,
  ):
      """Initializes a TorchPolicy instance.
  
      Args:
          observation_space: Observation space of the policy.
          action_space: Action space of the policy.
          config: The Policy's config dict.
          model: PyTorch policy module. Given observations as
              input, this module must return a list of outputs where the
              first item is action logits, and the rest can be any value.
          loss: Callable that returns one or more (a list of) scalar loss
              terms.
          action_distribution_class: Class for a torch action distribution.
          action_sampler_fn: A callable returning either a sampled action,
              its log-likelihood and updated state or a sampled action, its
              log-likelihood, updated state and action distribution inputs
              given Policy, ModelV2, input_dict, state batches (optional),
              explore, and timestep. Provide `action_sampler_fn` if you would
              like to have full control over the action computation step,
              including the model forward pass, possible sampling from a
              distribution, and exploration logic.
              Note: If `action_sampler_fn` is given, `action_distribution_fn`
              must be None. If both `action_sampler_fn` and
              `action_distribution_fn` are None, RLlib will simply pass
              inputs through `self.model` to get distribution inputs, create
              the distribution object, sample from it, and apply some
              exploration logic to the results.
              The callable takes as inputs: Policy, ModelV2, input_dict
              (SampleBatch), state_batches (optional), explore, and timestep.
          action_distribution_fn: A callable returning distribution inputs
              (parameters), a dist-class to generate an action distribution
              object from, and internal-state outputs (or an empty list if
              not applicable).
              Provide `action_distribution_fn` if you would like to only
              customize the model forward pass call. The resulting
              distribution parameters are then used by RLlib to create a
              distribution object, sample from it, and execute any
              exploration logic.
              Note: If `action_distribution_fn` is given, `action_sampler_fn`
              must be None. If both `action_sampler_fn` and
              `action_distribution_fn` are None, RLlib will simply pass
              inputs through `self.model` to get distribution inputs, create
              the distribution object, sample from it, and apply some
              exploration logic to the results.
              The callable takes as inputs: Policy, ModelV2, ModelInputDict,
              explore, timestep, is_training.
          max_seq_len: Max sequence length for LSTM training.
          get_batch_divisibility_req: Optional callable that returns the
              divisibility requirement for sample batches given the Policy.
      """
      self.framework = config["framework"] = "torch"
      self._loss_initialized = False
      super().__init__(observation_space, action_space, config)
  
      # Create multi-GPU model towers, if necessary.
      # - The central main model will be stored under self.model, residing
      #   on self.device (normally, a CPU).
      # - Each GPU will have a copy of that model under
      #   self.model_gpu_towers, matching the devices in self.devices.
      # - Parallelization is done by splitting the train batch and passing
      #   it through the model copies in parallel, then averaging over the
      #   resulting gradients, applying these averages on the main model and
      #   updating all towers' weights from the main model.
      # - In case of just one device (1 (fake or real) GPU or 1 CPU), no
      #   parallelization will be done.
  
      # If no Model is provided, build a default one here.
      if model is None:
          dist_class, logit_dim = ModelCatalog.get_action_dist(
              action_space,
              self.config["model"],
              framework=self.framework,
          )
          model = ModelCatalog.get_model_v2(
              obs_space=self.observation_space,
              action_space=self.action_space,
              num_outputs=logit_dim,
              model_config=self.config["model"],
              framework=self.framework,
          )
          if action_distribution_class is None:
              action_distribution_class = dist_class
  
      # Get devices to build the graph on.
      num_gpus = self._get_num_gpus_for_policy()
      gpu_ids = list(range(torch.cuda.device_count()))
      logger.info(f"Found {len(gpu_ids)} visible cuda devices.")
  
      # Place on one or more CPU(s) when either:
      # - Fake GPU mode.
      # - num_gpus=0 (either set by user or we are in local_mode=True).
      # - No GPUs available.
      if config["_fake_gpus"] or num_gpus == 0 or not gpu_ids:
          self.device = torch.device("cpu")
          self.devices = [
              self.device for _ in range(int(math.ceil(num_gpus)) or 1)
          ]
          self.model_gpu_towers = [
              model if i == 0 else copy.deepcopy(model)
              for i in range(int(math.ceil(num_gpus)) or 1)
          ]
          if hasattr(self, "target_model"):
              self.target_models = {
                  m: self.target_model for m in self.model_gpu_towers
              }
          self.model = model
      # Place on one or more actual GPU(s), when:
      # - num_gpus > 0 (set by user) AND
      # - local_mode=False AND
      # - actual GPUs available AND
      # - non-fake GPU mode.
      else:
          # We are a remote worker (WORKER_MODE=1):
          # GPUs should be assigned to us by ray.
          if (
              ray._private.worker._mode()
              == ray._private.worker.WORKER_MODE
          ):
              gpu_ids = ray.get_gpu_ids()
  
          if len(gpu_ids) < num_gpus:
              raise ValueError(
                  "TorchPolicy was not able to find enough GPU IDs! Found "
                  f"{gpu_ids}, but num_gpus={num_gpus}."
              )
  
          self.devices = [
              torch.device("cuda:{}".format(i))
              for i, id_ in enumerate(gpu_ids)
              if i < num_gpus
          ]
          self.device = self.devices[0]
          ids = [id_ for i, id_ in enumerate(gpu_ids) if i < num_gpus]
          self.model_gpu_towers = []
          for i, _ in enumerate(ids):
              model_copy = copy.deepcopy(model)
              self.model_gpu_towers.append(
                  model_copy.to(self.devices[i])
              )
          if hasattr(self, "target_model"):
              self.target_models = {
                  m: copy.deepcopy(self.target_model).to(
                      self.devices[i]
                  )
                  for i, m in enumerate(self.model_gpu_towers)
              }
          self.model = self.model_gpu_towers[0]
  
      # Lock used for locking some methods on the object-level.
      # This prevents possible race conditions when calling the model
      # first, then its value function (e.g. in a loss function), in
      # between of which another model call is made (e.g. to compute an
      # action).
      self._lock = threading.RLock()
  
      self._state_inputs = self.model.get_initial_state()
      self._is_recurrent = len(self._state_inputs) > 0
      # Auto-update model's inference view requirements, if recurrent.
      self._update_model_view_requirements_from_init_state()
      # Combine view_requirements for Model and Policy.
      self.view_requirements.update(self.model.view_requirements)
  
      self.exploration = self._create_exploration()
      self.unwrapped_model = (
          model  # used to support DistributedDataParallel
      )
      # To ensure backward compatibility:
      # Old way: If `loss` provided here, use as-is (as a function).
      if loss is not None:
          self._loss = loss
      # New way: Convert the overridden `self.loss` into a plain function,
      # so it can be called the same way as `loss` would be, ensuring
      # backward compatibility.
      elif self.loss.__func__.__qualname__ != "Policy.loss":
          self._loss = self.loss.__func__
      # `loss` not provided nor overridden from Policy -> Set to None.
      else:
          self._loss = None
      self._optimizers = force_list(self.optimizer())
      # Store, which params (by index within the model's list of
      # parameters) should be updated per optimizer.
      # Maps optimizer idx to set or param indices.
      self.multi_gpu_param_groups: List[Set[int]] = []
      main_params = {
          p: i for i, p in enumerate(self.model.parameters())
      }
      for o in self._optimizers:
          param_indices = []
          for pg_idx, pg in enumerate(o.param_groups):
              for p in pg["params"]:
                  param_indices.append(main_params[p])
          self.multi_gpu_param_groups.append(set(param_indices))
  
      # Create n sample-batch buffers (num_multi_gpu_tower_stacks), each
      # one with m towers (num_gpus).
      num_buffers = self.config.get("num_multi_gpu_tower_stacks", 1)
      self._loaded_batches = [[] for _ in range(num_buffers)]
  
      self.dist_class = action_distribution_class
      self.action_sampler_fn = action_sampler_fn
      self.action_distribution_fn = action_distribution_fn
  
      # If set, means we are using distributed allreduce during learning.
      self.distributed_world_size = None
  
      self.max_seq_len = max_seq_len
      self.batch_divisibility_req = (
          get_batch_divisibility_req(self)
          if callable(get_batch_divisibility_req)
          else (get_batch_divisibility_req or 1)
      )
  
imports: import copy
  import functools
  import logging
  import math
  import os
  import threading
  import time
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
  import gymnasium
  import numpy
  import tree
  import ray
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.policy import Policy, PolicyState
  from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils import NullContextManager, force_list
  from ray.rllib.utils.annotations import DeveloperAPI, override
  from ray.rllib.utils.error import ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics import DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, NUM_AGENT_STEPS_TRAINED, NUM_GRAD_UPDATES_LIFETIME
  from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.spaces.space_utils import normalize_action
  from ray.rllib.utils.threading import with_lock
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
  from ray.rllib.utils.typing import AlgorithmConfigDict, GradInfoDict, ModelGradients, ModelWeights, TensorStructType, TensorType
label_desc: Use image node_id 1 to create a new TorchPolicy object from inherited base classes: Policy with example: obj = TorchPolicy(observation_space, action_space, config)