class_docstr: PyTorch specific Policy class to use with RLlib.
class_docstr_tok: ["PyTorch","specific","Policy","class","to","use","with","RLlib","."]
docstr: Exports the Policy's Model to local directory for serving.
  
  Creates a TorchScript model and saves it.
  
  Args:
      export_dir: Local writable directory or filename.
      onnx: If given, will export model in ONNX format. The
          value of this parameter set the ONNX OpSet version to use.
docstr_tok: ["Exports","the","Policy","'s","Model","to","local","directory","for","serving",".","Creates","a","TorchScript","model","and","saves","it",".","Args",":","export_dir",":","Local","writable","directory","or","filename",".","onnx",":","If","given",",","will","export","model","in","ONNX","format",".","The","value","of","this","parameter","set","the","ONNX","OpSet","version","to","use","."]
code: def export_model(
      self, export_dir: str, onnx: Optional[int] = None
  ) -> None:
      """Exports the Policy's Model to local directory for serving.
  
      Creates a TorchScript model and saves it.
  
      Args:
          export_dir: Local writable directory or filename.
          onnx: If given, will export model in ONNX format. The
              value of this parameter set the ONNX OpSet version to use.
      """
      os.makedirs(export_dir, exist_ok=True)
  
      if onnx:
          self._lazy_tensor_dict(self._dummy_batch)
          # Provide dummy state inputs if not an RNN (torch cannot jit with
          # returned empty internal states list).
          if "state_in_0" not in self._dummy_batch:
              self._dummy_batch["state_in_0"] = self._dummy_batch[
                  SampleBatch.SEQ_LENS
              ] = np.array([1.0])
          seq_lens = self._dummy_batch[SampleBatch.SEQ_LENS]
  
          state_ins = []
          i = 0
          while "state_in_{}".format(i) in self._dummy_batch:
              state_ins.append(
                  self._dummy_batch["state_in_{}".format(i)]
              )
              i += 1
          dummy_inputs = {
              k: self._dummy_batch[k]
              for k in self._dummy_batch.keys()
              if k != "is_training"
          }
  
          file_name = os.path.join(export_dir, "model.onnx")
          torch.onnx.export(
              self.model,
              (dummy_inputs, state_ins, seq_lens),
              file_name,
              export_params=True,
              opset_version=onnx,
              do_constant_folding=True,
              input_names=list(dummy_inputs.keys())
              + ["state_ins", SampleBatch.SEQ_LENS],
              output_names=["output", "state_outs"],
              dynamic_axes={
                  k: {0: "batch_size"}
                  for k in list(dummy_inputs.keys())
                  + ["state_ins", SampleBatch.SEQ_LENS]
              },
          )
      # Save the torch.Model (architecture and weights, so it can be retrieved
      # w/o access to the original (custom) Model or Policy code).
      else:
          filename = os.path.join(export_dir, "model.pt")
          try:
              torch.save(self.model, f=filename)
          except Exception:
              if os.path.exists(filename):
                  os.remove(filename)
              logger.warning(ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL)
  
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 24 for calling the TorchPolicy obj's underlying member method code with example usage: obj.export_model(export_dir, onnx) without return types