class_docstr: PyTorch specific Policy class to use with RLlib.
class_docstr_tok: ["PyTorch","specific","Policy","class","to","use","with","RLlib","."]
code: def learn_on_loaded_batch(
      self, offset: int = 0, buffer_index: int = 0
  ):
      if not self._loaded_batches[buffer_index]:
          raise ValueError(
              "Must call Policy.load_batch_into_buffer() before "
              "Policy.learn_on_loaded_batch()!"
          )
  
      # Get the correct slice of the already loaded batch to use,
      # based on offset and batch size.
      device_batch_size = self.config.get(
          "sgd_minibatch_size", self.config["train_batch_size"]
      ) // len(self.devices)
  
      # Set Model to train mode.
      if self.model_gpu_towers:
          for t in self.model_gpu_towers:
              t.train()
  
      # Shortcut for 1 CPU only: Batch should already be stored in
      # `self._loaded_batches`.
      if len(self.devices) == 1 and self.devices[0].type == "cpu":
          assert buffer_index == 0
          if device_batch_size >= len(self._loaded_batches[0][0]):
              batch = self._loaded_batches[0][0]
          else:
              batch = self._loaded_batches[0][0][
                  offset : offset + device_batch_size
              ]
          return self.learn_on_batch(batch)
  
      if len(self.devices) > 1:
          # Copy weights of main model (tower-0) to all other towers.
          state_dict = self.model.state_dict()
          # Just making sure tower-0 is really the same as self.model.
          assert self.model_gpu_towers[0] is self.model
          for tower in self.model_gpu_towers[1:]:
              tower.load_state_dict(state_dict)
  
      if device_batch_size >= sum(
          len(s) for s in self._loaded_batches[buffer_index]
      ):
          device_batches = self._loaded_batches[buffer_index]
      else:
          device_batches = [
              b[offset : offset + device_batch_size]
              for b in self._loaded_batches[buffer_index]
          ]
  
      # Callback handling.
      batch_fetches = {}
      for i, batch in enumerate(device_batches):
          custom_metrics = {}
          self.callbacks.on_learn_on_batch(
              policy=self, train_batch=batch, result=custom_metrics
          )
          batch_fetches[f"tower_{i}"] = {
              "custom_metrics": custom_metrics
          }
  
      # Do the (maybe parallelized) gradient calculation step.
      tower_outputs = self._multi_gpu_parallel_grad_calc(device_batches)
  
      # Mean-reduce gradients over GPU-towers (do this on CPU: self.device).
      all_grads = []
      for i in range(len(tower_outputs[0][0])):
          if tower_outputs[0][0][i] is not None:
              all_grads.append(
                  torch.mean(
                      torch.stack(
                          [
                              t[0][i].to(self.device)
                              for t in tower_outputs
                          ]
                      ),
                      dim=0,
                  )
              )
          else:
              all_grads.append(None)
      # Set main model's grads to mean-reduced values.
      for i, p in enumerate(self.model.parameters()):
          p.grad = all_grads[i]
  
      self.apply_gradients(_directStepOptimizerSingleton)
  
      self.num_grad_updates += 1
  
      for i, (model, batch) in enumerate(
          zip(self.model_gpu_towers, device_batches)
      ):
          batch_fetches[f"tower_{i}"].update(
              {
                  LEARNER_STATS_KEY: self.extra_grad_info(batch),
                  "model": model.metrics(),
                  NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates,
                  # -1, b/c we have to measure this diff before we do the update
                  # above.
                  DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: (
                      self.num_grad_updates
                      - 1
                      - (batch.num_grad_updates or 0)
                  ),
              }
          )
      batch_fetches.update(self.extra_compute_grad_fetches())
  
      return batch_fetches
  
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 8 for calling the TorchPolicy obj's underlying member method code with example usage: obj.learn_on_loaded_batch(offset, buffer_index) and returns: batch_fetches, self