class_docstr: Abstract base class for TF-based RL and Bandits agents.
  
  The agent serves the following purposes:
  
  * Training by reading minibatches of `experience`, and updating some set
    of network weights (using the `train` method).
  
  * Exposing `policy` objects which can be used to interact with an environment:
    either to explore and collect new training data, or to maximize reward
    in the given task.
  
  The agents' main training methods and properties are:
  
  * `initialize`: Perform any self-initialization before training.
  
  * `train`: This method reads minibatch experience from a replay buffer or
    logs on disk, and updates some internal networks.
  
  * `preprocess_sequence`: Some algorithms need to perform sequence
    preprocessing on logs containing "full episode" or "long subset" sequences,
    to create intermediate items that can then be used by `train`, even if
    `train` does not see the full sequences.  In many cases this is just
    the identity: it passes experience through untouched.  This function
    is typically passed to the argument
  
    `ReplayBuffer.as_dataset(..., sequence_preprocess_fn=...)`
  
  * `training_data_spec`: Property that describes the structure expected of
    the `experience` argument passed to `train`.
  
  * `train_sequence_length`: Property that describes the **second** dimension
    of all tensors in the `experience` argument passed `train`.  All tensors
    passed to train must have the shape `[batch_size, sequence_length, ...]`,
    and some Agents require this to be a fixed value.  For example, in regular
    `DQN`, this second `sequence_length` dimension must be equal to `2` in all
    `experience`.  In contrast, `n-step DQN` will have this equal to `n + 1` and
    `DQN` agents constructed with `RNN` networks will have this equal to `None`,
    meaning any length sequences are allowed.
  
    This value may be `None`, to mean minibatches containing subsequences of any
    length are allowed (so long as they're all the same length).  This is
    typically the case with agents constructed with `RNN` networks.
  
    This value is typically passed as a ReplayBuffer's
    `as_dataset(..., num_steps=...)` argument.
  
  * `collect_data_spec`: Property that describes the structure expected of
    experience collected by `agent.collect_policy`.  This is typically
    identical to `training_data_spec`, but may be different if
    `preprocess_sequence` method is not the identity.  In this case,
    `preprocess_sequence` is expected to read sequences matching
    `collect_data_spec` and emit sequences matching `training_data_spec`.
  
  The agent exposes `TFPolicy` objects for interacting with environments:
  
  * `policy`: Property that returns a policy meant for "exploiting" the
    environment to its best ability.  This tends to mean the "production" policy
    that doesn't collect additional info for training.  Works best when
    the agent is fully trained.
  
    TODO(b/154870654): Not all agents are properly exporting properly greedy
    "production" policies yet.  We have to clean this up.  In particular,
    we have to update PPO and SAC's `policy` objects.
  
  * `collect_policy`: Property that returns a policy meant for "exploring"
    the environment to collect more data for training.  This tends to mean
    a policy involves some level of randomized behavior and additional info
    logging.
  
  * `time_step_spec`: Property describing the observation and reward signatures
    of the environment this agent's policies operate in.
  
  * `action_spec`: Property describing the action signatures of the environment
    this agent's policies operate in.
  
  
  **NOTE**: For API consistency, subclasses are not allowed to override public
  methods of `TFAgent` class. Instead, they may implement the protected methods
  including `_initialize`, `_train`, and `_preprocess_sequence`. This
  public-calls-private convention allowed this base class to do things like
  properly add `spec` and shape checks, which provide users an easier experience
  when debugging their environments and networks.
docstr: Meant to be called by subclass constructors.
  
  Args:
    time_step_spec: A nest of tf.TypeSpec representing the time_steps.
      Provided by the user.
    action_spec: A nest of BoundedTensorSpec representing the actions.
      Provided by the user.
    policy: An instance of `tf_policy.TFPolicy` representing the Agent's
      current policy.
    collect_policy: An instance of `tf_policy.TFPolicy` representing the
      Agent's current data collection policy (used to set `self.step_spec`).
    train_sequence_length: A python integer or `None`, signifying the number
      of time steps required from tensors in `experience` as passed to
      `train()`.  All tensors in `experience` will be shaped `[B, T, ...]` but
      for certain agents, `T` should be fixed.  For example, DQN requires
      transitions in the form of 2 time steps, so for a non-RNN DQN Agent, set
      this value to 2.  For agents that don't care, or which can handle `T`
      unknown at graph build time (i.e. most RNN-based agents), set this
      argument to `None`.
    num_outer_dims: The number of outer dimensions for the agent. Must be
      either 1 or 2. If 2, training will require both a batch_size and time
      dimension on every Tensor; if 1, training will require only a batch_size
      outer dimension.
    training_data_spec: A nest of TensorSpec specifying the structure of data
      the train() function expects. If None, defaults to the trajectory_spec
      of the collect_policy.
    debug_summaries: A bool; if true, subclasses should gather debug
      summaries.
    summarize_grads_and_vars: A bool; if true, subclasses should additionally
      collect gradient and variable summaries.
    enable_summaries: A bool; if false, subclasses should not gather any
      summaries (debug or otherwise); subclasses should gate *all* summaries
      using either `summaries_enabled`, `debug_summaries`, or
      `summarize_grads_and_vars` properties.
    train_step_counter: An optional counter to increment every time the train
      op is run.  Defaults to the global_step.
  
  Raises:
    ValueError: If `num_outer_dims` is not in `[1, 2]`.
docstr_tok: ["Meant","to","be","called","by","subclass","constructors",".","Args",":","time_step_spec",":","A","nest","of","tf.TypeSpec","representing","the","time_steps",".","Provided","by","the","user",".","action_spec",":","A","nest","of","BoundedTensorSpec","representing","the","actions",".","Provided","by","the","user",".","policy",":","An","instance","of","`","tf_policy.TFPolicy","`","representing","the","Agent's","current","policy",".","collect_policy",":","An","instance","of","`","tf_policy.TFPolicy","`","representing","the","Agent","'s","current","data","collection","policy","(","used","to","set","`","self.step_spec","`",")",".","train_sequence_length",":","A","python","integer","or","`","None","`",",","signifying","the","number","of","time","steps","required","from","tensors","in","`","experience","`","as","passed","to","`","train","(",")","`",".","All","tensors","in","`","experience","`","will","be","shaped","`","[","B",",","T",",","...","]","`","but","for","certain","agents",",","`","T","`","should","be","fixed",".","For","example",",","DQN","requires","transitions","in","the","form","of","2","time","steps",",","so","for","a","non-RNN","DQN","Agent",",","set","this","value","to","2",".","For","agents","that","do","n't","care",",","or","which","can","handle","`","T","`","unknown","at","graph","build","time","(","i.e",".","most","RNN-based","agents",")",",","set","this","argument","to","`","None","`",".","num_outer_dims",":","The","number","of","outer","dimensions","for","the","agent",".","Must","be","either","1","or","2",".","If","2",",","training","will","require","both","a","batch_size","and","time","dimension","on","every","Tensor",";","if","1",",","training","will","require","only","a","batch_size","outer","dimension",".","training_data_spec",":","A","nest","of","TensorSpec","specifying","the","structure","of","data","the","train","(",")","function","expects",".","If","None",",","defaults","to","the","trajectory_spec","of","the","collect_policy",".","debug_summaries",":","A","bool",";","if","true",",","subclasses","should","gather","debug","summaries",".","summarize_grads_and_vars",":","A","bool",";","if","true",",","subclasses","should","additionally","collect","gradient","and","variable","summaries",".","enable_summaries",":","A","bool",";","if","false",",","subclasses","should","not","gather","any","summaries","(","debug","or","otherwise",")",";","subclasses","should","gate","*","all","*","summaries","using","either","`","summaries_enabled","`",",","`","debug_summaries","`",",","or","`","summarize_grads_and_vars","`","properties",".","train_step_counter",":","An","optional","counter","to","increment","every","time","the","train","op","is","run",".","Defaults","to","the","global_step",".","Raises",":","ValueError",":","If","`","num_outer_dims","`","is","not","in","`","[","1",",","2","]","`","."]
code: def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      policy: tf_policy.TFPolicy,
      collect_policy: tf_policy.TFPolicy,
      train_sequence_length: Optional[int],
      num_outer_dims: int = 2,
      training_data_spec: Optional[types.NestedTensorSpec] = None,
      debug_summaries: bool = False,
      summarize_grads_and_vars: bool = False,
      enable_summaries: bool = True,
      train_step_counter: Optional[tf.Variable] = None,
  ):
      """Meant to be called by subclass constructors.
  
      Args:
        time_step_spec: A nest of tf.TypeSpec representing the time_steps.
          Provided by the user.
        action_spec: A nest of BoundedTensorSpec representing the actions.
          Provided by the user.
        policy: An instance of `tf_policy.TFPolicy` representing the Agent's
          current policy.
        collect_policy: An instance of `tf_policy.TFPolicy` representing the
          Agent's current data collection policy (used to set `self.step_spec`).
        train_sequence_length: A python integer or `None`, signifying the number
          of time steps required from tensors in `experience` as passed to
          `train()`.  All tensors in `experience` will be shaped `[B, T, ...]` but
          for certain agents, `T` should be fixed.  For example, DQN requires
          transitions in the form of 2 time steps, so for a non-RNN DQN Agent, set
          this value to 2.  For agents that don't care, or which can handle `T`
          unknown at graph build time (i.e. most RNN-based agents), set this
          argument to `None`.
        num_outer_dims: The number of outer dimensions for the agent. Must be
          either 1 or 2. If 2, training will require both a batch_size and time
          dimension on every Tensor; if 1, training will require only a batch_size
          outer dimension.
        training_data_spec: A nest of TensorSpec specifying the structure of data
          the train() function expects. If None, defaults to the trajectory_spec
          of the collect_policy.
        debug_summaries: A bool; if true, subclasses should gather debug
          summaries.
        summarize_grads_and_vars: A bool; if true, subclasses should additionally
          collect gradient and variable summaries.
        enable_summaries: A bool; if false, subclasses should not gather any
          summaries (debug or otherwise); subclasses should gate *all* summaries
          using either `summaries_enabled`, `debug_summaries`, or
          `summarize_grads_and_vars` properties.
        train_step_counter: An optional counter to increment every time the train
          op is run.  Defaults to the global_step.
  
      Raises:
        ValueError: If `num_outer_dims` is not in `[1, 2]`.
      """
      common.check_tf1_allowed()
      common.tf_agents_gauge.get_cell("TFAgent").set(True)
      common.tf_agents_gauge.get_cell(str(type(self))).set(True)
      if not isinstance(time_step_spec, ts.TimeStep):
          raise TypeError(
              "The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.".format(
                  type(time_step_spec)
              )
          )
  
      if num_outer_dims not in [1, 2]:
          raise ValueError("num_outer_dims must be in [1, 2].")
  
      time_step_spec = tensor_spec.from_spec(time_step_spec)
      action_spec = tensor_spec.from_spec(action_spec)
      self._time_step_spec = time_step_spec
      self._action_spec = action_spec
      self._policy = policy
      self._collect_policy = collect_policy
      self._train_sequence_length = train_sequence_length
      self._num_outer_dims = num_outer_dims
      self._debug_summaries = debug_summaries
      self._summarize_grads_and_vars = summarize_grads_and_vars
      self._enable_summaries = enable_summaries
      self._training_data_spec = training_data_spec
      # Data context for data collected directly from the collect policy.
      self._collect_data_context = data_converter.DataContext(
          time_step_spec=self._time_step_spec,
          action_spec=self._action_spec,
          info_spec=collect_policy.info_spec,
      )
      # Data context for data passed to train().  May be different if
      # training_data_spec is provided.
      if training_data_spec is not None:
          training_data_spec = tensor_spec.from_spec(training_data_spec)
          # training_data_spec can be anything; so build a data_context
          # via best-effort with fall-backs to the collect data spec.
          training_discount_spec = getattr(
              training_data_spec, "discount", time_step_spec.discount
          )
          training_observation_spec = getattr(
              training_data_spec,
              "observation",
              time_step_spec.observation,
          )
          training_reward_spec = getattr(
              training_data_spec, "reward", time_step_spec.reward
          )
          training_step_type_spec = getattr(
              training_data_spec, "step_type", time_step_spec.step_type
          )
          training_policy_info_spec = getattr(
              training_data_spec,
              "policy_info",
              collect_policy.info_spec,
          )
          training_action_spec = getattr(
              training_data_spec, "action", action_spec
          )
          self._data_context = data_converter.DataContext(
              time_step_spec=ts.TimeStep(
                  discount=training_discount_spec,
                  observation=training_observation_spec,
                  reward=training_reward_spec,
                  step_type=training_step_type_spec,
              ),
              action_spec=training_action_spec,
              info_spec=training_policy_info_spec,
          )
      else:
          self._data_context = data_converter.DataContext(
              time_step_spec=time_step_spec,
              action_spec=action_spec,
              info_spec=collect_policy.info_spec,
          )
      if train_step_counter is None:
          train_step_counter = (
              tf.compat.v1.train.get_or_create_global_step()
          )
      self._train_step_counter = train_step_counter
      self._train_fn = common.function_in_tf1()(self._train)
      self._initialize_fn = common.function_in_tf1()(self._initialize)
      self._preprocess_sequence_fn = common.function_in_tf1()(
          self._preprocess_sequence
      )
      self._loss_fn = common.function_in_tf1()(self._loss)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  import collections
  from typing import Optional
  import six
  import tensorflow
  from tf_agents.agents import data_converter
  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 eager_utils
label_desc: Use image node_id 1 to create a new TFAgent object from inherited base classes: tf with example: obj = TFAgent(time_step_spec, action_spec, policy, collect_policy, train_sequence_length, num_outer_dims, training_data_spec, debug_summaries, summarize_grads_and_vars, enable_summaries, train_step_counter)