class_docstr: Exposes a numpy API for saved_model policies in Eager mode.
class_docstr_tok: ["Exposes","a","numpy","API","for","saved_model","policies","in","Eager","mode","."]
docstr: Initializes a PyPolicy from a saved_model.
  
  *Note* (b/151318119): BoundedSpecs are converted to regular specs when saved
  into a proto as the `nested_structure_coder` from TF currently doesn't
  handle BoundedSpecs. Shape and dtypes will still match the original specs.
  
  Args:
    model_path: Path to a saved_model generated by the `policy_saver`.
    time_step_spec: Optional nested structure of ArraySpecs describing the
      policy's `time_step_spec`. This is not used by the
      SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
      part of the public policy API.
    action_spec: Optional nested structure of `ArraySpecs` describing the
      policy's `action_spec`. This is not used by the
      SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
      part of the public policy API.
    policy_state_spec: Optional nested structure of `ArraySpecs` describing
      the policy's `policy_state_spec`. This is not used by the
      SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
      part of the public policy API.
    info_spec: Optional nested structure of `ArraySpecs` describing the
      policy's `info_spec`. This is not used by the SavedModelPyTFEagerPolicy,
      but may be accessed by other objects as it is part of the public policy
      API.
    load_specs_from_pbtxt: If True the specs will be loaded from the proto
      file generated by the `policy_saver`.
    use_tf_function: See PyTFEagerPolicyBase.
    batch_time_steps: See PyTFEagerPolicyBase.
docstr_tok: ["Initializes","a","PyPolicy","from","a","saved_model",".","*","Note","*","(","b\/151318119",")",":","BoundedSpecs","are","converted","to","regular","specs","when","saved","into","a","proto","as","the","`","nested_structure_coder","`","from","TF","currently","doesn't","handle","BoundedSpecs",".","Shape","and","dtypes","will","still","match","the","original","specs",".","Args",":","model_path",":","Path","to","a","saved_model","generated","by","the","`","policy_saver","`",".","time_step_spec",":","Optional","nested","structure","of","ArraySpecs","describing","the","policy","'s","`","time_step_spec","`",".","This","is","not","used","by","the","SavedModelPyTFEagerPolicy",",","but","may","be","accessed","by","other","objects","as","it","is","part","of","the","public","policy","API",".","action_spec",":","Optional","nested","structure","of","`","ArraySpecs","`","describing","the","policy","'s","`","action_spec","`",".","This","is","not","used","by","the","SavedModelPyTFEagerPolicy",",","but","may","be","accessed","by","other","objects","as","it","is","part","of","the","public","policy","API",".","policy_state_spec",":","Optional","nested","structure","of","`","ArraySpecs","`","describing","the","policy","'s","`","policy_state_spec","`",".","This","is","not","used","by","the","SavedModelPyTFEagerPolicy",",","but","may","be","accessed","by","other","objects","as","it","is","part","of","the","public","policy","API",".","info_spec",":","Optional","nested","structure","of","`","ArraySpecs","`","describing","the","policy","'s","`","info_spec","`",".","This","is","not","used","by","the","SavedModelPyTFEagerPolicy",",","but","may","be","accessed","by","other","objects","as","it","is","part","of","the","public","policy","API",".","load_specs_from_pbtxt",":","If","True","the","specs","will","be","loaded","from","the","proto","file","generated","by","the","`","policy_saver","`",".","use_tf_function",":","See","PyTFEagerPolicyBase",".","batch_time_steps",":","See","PyTFEagerPolicyBase","."]
code: def __init__(
      self,
      model_path: Text,
      time_step_spec: Optional[ts.TimeStep] = None,
      action_spec: Optional[types.NestedTensorSpec] = None,
      policy_state_spec: types.NestedTensorSpec = (),
      info_spec: types.NestedTensorSpec = (),
      load_specs_from_pbtxt: bool = False,
      use_tf_function: bool = False,
      batch_time_steps=True,
  ):
      """Initializes a PyPolicy from a saved_model.
  
      *Note* (b/151318119): BoundedSpecs are converted to regular specs when saved
      into a proto as the `nested_structure_coder` from TF currently doesn't
      handle BoundedSpecs. Shape and dtypes will still match the original specs.
  
      Args:
        model_path: Path to a saved_model generated by the `policy_saver`.
        time_step_spec: Optional nested structure of ArraySpecs describing the
          policy's `time_step_spec`. This is not used by the
          SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
          part of the public policy API.
        action_spec: Optional nested structure of `ArraySpecs` describing the
          policy's `action_spec`. This is not used by the
          SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
          part of the public policy API.
        policy_state_spec: Optional nested structure of `ArraySpecs` describing
          the policy's `policy_state_spec`. This is not used by the
          SavedModelPyTFEagerPolicy, but may be accessed by other objects as it is
          part of the public policy API.
        info_spec: Optional nested structure of `ArraySpecs` describing the
          policy's `info_spec`. This is not used by the SavedModelPyTFEagerPolicy,
          but may be accessed by other objects as it is part of the public policy
          API.
        load_specs_from_pbtxt: If True the specs will be loaded from the proto
          file generated by the `policy_saver`.
        use_tf_function: See PyTFEagerPolicyBase.
        batch_time_steps: See PyTFEagerPolicyBase.
      """
      policy = tf.compat.v2.saved_model.load(model_path)
      self._checkpoint = tf.train.Checkpoint(policy=policy)
      if not (time_step_spec or load_specs_from_pbtxt):
          raise ValueError(
              "To load a SavedModel policy you have to provide the specs, or"
              "enable loading from proto."
          )
      policy_specs = None
      if not time_step_spec and load_specs_from_pbtxt:
          spec_path = os.path.join(
              model_path, policy_saver.POLICY_SPECS_PBTXT
          )
          policy_specs = policy_saver.specs_from_collect_data_spec(
              tensor_spec.from_pbtxt_file(spec_path)
          )
          time_step_spec = policy_specs["time_step_spec"]
          action_spec = policy_specs["action_spec"]
          policy_state_spec = policy_specs["policy_state_spec"]
          info_spec = policy_specs["info_spec"]
      super(SavedModelPyTFEagerPolicy, self).__init__(
          policy,
          time_step_spec,
          action_spec,
          policy_state_spec,
          info_spec,
          use_tf_function,
          batch_time_steps,
      )
      # Override collect data_spec with whatever was loaded instead of relying
      # on trajectory_data_spec.
      if policy_specs:
          self._collect_data_spec = policy_specs["collect_data_spec"]
  
      # Call action function of the saved_model policy w/ dummy data to
      # initialize the graph. Otherwise the first call to the action function may
      # be 10x slower than the rest of the calls on GPU.
      dummy_time_step = tf.nest.map_structure(
          lambda s: tf.zeros(s.shape, s.dtype), time_step_spec
      )
      dummy_time_step = nest_utils.batch_nested_tensors(dummy_time_step)
      dummy_policy_state = tf.nest.map_structure(
          lambda s: tf.zeros(s.shape, s.dtype), policy_state_spec
      )
      dummy_policy_state = nest_utils.batch_nested_tensors(
          dummy_policy_state
      )
      self._policy_action_fn(dummy_time_step, dummy_policy_state)
      self._train_step_from_last_restored_checkpoint_path = None
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from typing import Optional, Text
  from absl import logging
  import gin
  import tensorflow
  from tf_agents.policies import policy_saver
  from tf_agents.policies import py_policy
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 to create a new SavedModelPyTFEagerPolicy object from inherited base classes: PyTFEagerPolicyBase with example: obj = SavedModelPyTFEagerPolicy(model_path, time_step_spec, action_spec, policy_state_spec, info_spec, load_specs_from_pbtxt, use_tf_function, batch_time_steps)