class_docstr: PyTorch specific Policy class to use with RLlib.
class_docstr_tok: ["PyTorch","specific","Policy","class","to","use","with","RLlib","."]
code: def set_state(self, state: PolicyState) -> None:
      # Set optimizer vars first.
      optimizer_vars = state.get("_optimizer_variables", None)
      if optimizer_vars:
          assert len(optimizer_vars) == len(self._optimizers)
          for o, s in zip(self._optimizers, optimizer_vars):
              # Torch optimizer param_groups include things like beta, etc. These
              # parameters should be left as scalar and not converted to tensors.
              # otherwise, torch.optim.step() will start to complain.
              optim_state_dict = {"param_groups": s["param_groups"]}
              optim_state_dict["state"] = convert_to_torch_tensor(
                  s["state"], device=self.device
              )
              o.load_state_dict(optim_state_dict)
      # Set exploration's state.
      if hasattr(self, "exploration") and "_exploration_state" in state:
          self.exploration.set_state(state=state["_exploration_state"])
  
      # Restore global timestep.
      self.global_timestep = state["global_timestep"]
  
      # Then the Policy's (NN) weights and connectors.
      super().set_state(state)
  
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 18 for calling the TorchPolicy obj's underlying member method code with example usage: obj.set_state(state) without return types