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_batch(
      self, postprocessed_batch: SampleBatch
  ) -> Dict[str, TensorType]:
      # Set Model to train mode.
      if self.model:
          self.model.train()
      # Callback handling.
      learn_stats = {}
      self.callbacks.on_learn_on_batch(
          policy=self,
          train_batch=postprocessed_batch,
          result=learn_stats,
      )
  
      # Compute gradients (will calculate all losses and `backward()`
      # them to get the grads).
      grads, fetches = self.compute_gradients(postprocessed_batch)
  
      # Step the optimizers.
      self.apply_gradients(_directStepOptimizerSingleton)
  
      self.num_grad_updates += 1
  
      if self.model:
          fetches["model"] = self.model.metrics()
  
      fetches.update(
          {
              "custom_metrics": learn_stats,
              NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count,
              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
                  - (postprocessed_batch.num_grad_updates or 0)
              ),
          }
      )
  
      return 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 5 for calling the TorchPolicy obj's underlying member method code with example usage: obj.learn_on_batch(postprocessed_batch) and returns: fetches