class_docstr: Returns samples with updated `policy_info` (a dictionary).
class_docstr_tok: ["Returns","samples","with","updated","`","policy_info","`","(","a","dictionary",")","."]
docstr: Builds a TFPolicy wrapping the given policy.
  
  PolicyInfoUpdaterWrapper class updates `policy_info` using a user-defined
  updater function. The main use case of this policy wrapper is to annotate
  `policy_info` with some auxiliary information. For example, appending
  an identifier to specify which model is used for current rollout.
  
  Args:
    policy: A policy implementing the tf_policy.TFPolicy interface.
    info_spec: User-defined `info_spec` which specifies the policy info after
      applying the updater function.
    updater_fn: An updater function that updates the `policy_info`. This is a
      callable that receives a `PolicyStep` and will return a dictionary of a
      tf.Tensor or sequence of tf.Tensor`s.  **NOTE** If `policy.distribution`
      is called, the `PolicyStep.action` object may contain a
      `tfp.distributions.Distribution` object instead of a `Tensor`.  The
      `updater_fn` must be able to handle both cases to be compatible with
      `PolicySaver`.
    name: The name of this policy. All variables in this module will fall
      under that name. Defaults to the class name.
docstr_tok: ["Builds","a","TFPolicy","wrapping","the","given","policy",".","PolicyInfoUpdaterWrapper","class","updates","`","policy_info","`","using","a","user-defined","updater","function",".","The","main","use","case","of","this","policy","wrapper","is","to","annotate","`","policy_info","`","with","some","auxiliary","information",".","For","example",",","appending","an","identifier","to","specify","which","model","is","used","for","current","rollout",".","Args",":","policy",":","A","policy","implementing","the","tf_policy.TFPolicy","interface",".","info_spec",":","User-defined","`","info_spec","`","which","specifies","the","policy","info","after","applying","the","updater","function",".","updater_fn",":","An","updater","function","that","updates","the","`","policy_info","`",".","This","is","a","callable","that","receives","a","`","PolicyStep","`","and","will","return","a","dictionary","of","a","tf.Tensor","or","sequence","of","tf.Tensor","`","s",".","*","*","NOTE","*","*","If","`","policy.distribution","`","is","called",",","the","`","PolicyStep.action","`","object","may","contain","a","`","tfp.distributions.Distribution","`","object","instead","of","a","`","Tensor","`",".","The","`","updater_fn","`","must","be","able","to","handle","both","cases","to","be","compatible","with","`","PolicySaver","`",".","name",":","The","name","of","this","policy",".","All","variables","in","this","module","will","fall","under","that","name",".","Defaults","to","the","class","name","."]
code: def __init__(
      self,
      policy: tf_policy.TFPolicy,
      info_spec: types.NestedTensorSpec,
      updater_fn: UpdaterFnType,
      name: Optional[Text] = None,
  ):
      """Builds a TFPolicy wrapping the given policy.
  
      PolicyInfoUpdaterWrapper class updates `policy_info` using a user-defined
      updater function. The main use case of this policy wrapper is to annotate
      `policy_info` with some auxiliary information. For example, appending
      an identifier to specify which model is used for current rollout.
  
      Args:
        policy: A policy implementing the tf_policy.TFPolicy interface.
        info_spec: User-defined `info_spec` which specifies the policy info after
          applying the updater function.
        updater_fn: An updater function that updates the `policy_info`. This is a
          callable that receives a `PolicyStep` and will return a dictionary of a
          tf.Tensor or sequence of tf.Tensor`s.  **NOTE** If `policy.distribution`
          is called, the `PolicyStep.action` object may contain a
          `tfp.distributions.Distribution` object instead of a `Tensor`.  The
          `updater_fn` must be able to handle both cases to be compatible with
          `PolicySaver`.
        name: The name of this policy. All variables in this module will fall
          under that name. Defaults to the class name.
      """
      super(PolicyInfoUpdaterWrapper, self).__init__(
          time_step_spec=policy.time_step_spec,
          action_spec=policy.action_spec,
          policy_state_spec=policy.policy_state_spec,
          info_spec=info_spec,
          emit_log_probability=policy.emit_log_probability,
          name=name,
      )
      self._wrapped_policy = policy
      self._info_spec = info_spec
      self._updater_fn = updater_fn
  
file_path: agents/tf_agents/policies/policy_info_updater_wrapper.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Callable, Dict, Optional, Sequence, Text, Union
  import tensorflow.compat.v2
  from tf_agents.policies import tf_policy
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
label_desc: Use image node_id 1 to create a new PolicyInfoUpdaterWrapper object from inherited base classes: tf_policy with example: obj = PolicyInfoUpdaterWrapper(policy, info_spec, updater_fn, name)