code: def train_eval(
      root_dir,
      env_name="HalfCheetah-v2",
      # Training params
      num_iterations=1600,
      actor_fc_layers=(64, 64),
      value_fc_layers=(64, 64),
      learning_rate=3e-4,
      collect_sequence_length=2048,
      minibatch_size=64,
      num_epochs=10,
      # Agent params
      importance_ratio_clipping=0.2,
      lambda_value=0.95,
      discount_factor=0.99,
      entropy_regularization=0.0,
      value_pred_loss_coef=0.5,
      use_gae=True,
      use_td_lambda_return=True,
      gradient_clipping=0.5,
      value_clipping=None,
      # Replay params
      reverb_port=None,
      replay_capacity=10000,
      # Others
      policy_save_interval=5000,
      summary_interval=1000,
      eval_interval=10000,
      eval_episodes=100,
      debug_summaries=False,
      summarize_grads_and_vars=False,
  ):
      """Trains and evaluates PPO (Importance Ratio Clipping).
  
      Args:
        root_dir: Main directory path where checkpoints, saved_models, and summaries
          will be written to.
        env_name: Name for the Mujoco environment to load.
        num_iterations: The number of iterations to perform collection and training.
        actor_fc_layers: List of fully_connected parameters for the actor network,
          where each item is the number of units in the layer.
        value_fc_layers: : List of fully_connected parameters for the value network,
          where each item is the number of units in the layer.
        learning_rate: Learning rate used on the Adam optimizer.
        collect_sequence_length: Number of steps to take in each collect run.
        minibatch_size: Number of elements in each mini batch. If `None`, the entire
          collected sequence will be treated as one batch.
        num_epochs: Number of iterations to repeat over all collected data per data
          collection step. (Schulman,2017) sets this to 10 for Mujoco, 15 for
          Roboschool and 3 for Atari.
        importance_ratio_clipping: Epsilon in clipped, surrogate PPO objective. For
          more detail, see explanation at the top of the doc.
        lambda_value: Lambda parameter for TD-lambda computation.
        discount_factor: Discount factor for return computation. Default to `0.99`
          which is the value used for all environments from (Schulman, 2017).
        entropy_regularization: Coefficient for entropy regularization loss term.
          Default to `0.0` because no entropy bonus was used in (Schulman, 2017).
        value_pred_loss_coef: Multiplier for value prediction loss to balance with
          policy gradient loss. Default to `0.5`, which was used for all
          environments in the OpenAI baseline implementation. This parameters is
          irrelevant unless you are sharing part of actor_net and value_net. In that
          case, you would want to tune this coeeficient, whose value depends on the
          network architecture of your choice.
        use_gae: If True (default False), uses generalized advantage estimation for
          computing per-timestep advantage. Else, just subtracts value predictions
          from empirical return.
        use_td_lambda_return: If True (default False), uses td_lambda_return for
          training value function; here: `td_lambda_return = gae_advantage +
          value_predictions`. `use_gae` must be set to `True` as well to enable TD
          -lambda returns. If `use_td_lambda_return` is set to True while `use_gae`
          is False, the empirical return will be used and a warning will be logged.
        gradient_clipping: Norm length to clip gradients.
        value_clipping: Difference between new and old value predictions are clipped
          to this threshold. Value clipping could be helpful when training very deep
          networks. Default: no clipping.
        reverb_port: Port for reverb server, if None, use a randomly chosen unused
          port.
        replay_capacity: The maximum number of elements for the replay buffer. Items
          will be wasted if this is smalled than collect_sequence_length.
        policy_save_interval: How often, in train_steps, the policy will be saved.
        summary_interval: How often to write data into Tensorboard.
        eval_interval: How often to run evaluation, in train_steps.
        eval_episodes: Number of episodes to evaluate over.
        debug_summaries: Boolean for whether to gather debug summaries.
        summarize_grads_and_vars: If true, gradient summaries will be written.
      """
      collect_env = suite_mujoco.load(env_name)
      eval_env = suite_mujoco.load(env_name)
      num_environments = 1
  
      (
          observation_tensor_spec,
          action_tensor_spec,
          time_step_tensor_spec,
      ) = spec_utils.get_tensor_specs(collect_env)
      # TODO(b/172267869): Remove this conversion once TensorNormalizer stops
      # converting float64 inputs to float32.
      observation_tensor_spec = tf.TensorSpec(
          dtype=tf.float32, shape=observation_tensor_spec.shape
      )
  
      train_step = train_utils.create_train_step()
      actor_net_builder = ppo_actor_network.PPOActorNetwork()
      actor_net = actor_net_builder.create_sequential_actor_net(
          actor_fc_layers, action_tensor_spec
      )
      value_net = value_network.ValueNetwork(
          observation_tensor_spec,
          fc_layer_params=value_fc_layers,
          kernel_initializer=tf.keras.initializers.Orthogonal(),
      )
  
      current_iteration = tf.Variable(0, dtype=tf.int64)
  
      def learning_rate_fn():
          # Linearly decay the learning rate.
          return learning_rate * (
              1 - current_iteration / num_iterations
          )
  
      agent = ppo_clip_agent.PPOClipAgent(
          time_step_tensor_spec,
          action_tensor_spec,
          optimizer=tf.keras.optimizers.legacy.Adam(
              learning_rate=learning_rate_fn, epsilon=1e-5
          ),
          actor_net=actor_net,
          value_net=value_net,
          importance_ratio_clipping=importance_ratio_clipping,
          lambda_value=lambda_value,
          discount_factor=discount_factor,
          entropy_regularization=entropy_regularization,
          value_pred_loss_coef=value_pred_loss_coef,
          # This is a legacy argument for the number of times we repeat the data
          # inside of the train function, incompatible with mini batch learning.
          # We set the epoch number from the replay buffer and tf.Data instead.
          num_epochs=1,
          use_gae=use_gae,
          use_td_lambda_return=use_td_lambda_return,
          gradient_clipping=gradient_clipping,
          value_clipping=value_clipping,
          # TODO(b/150244758): Default compute_value_and_advantage_in_train to False
          # after Reverb open source.
          compute_value_and_advantage_in_train=False,
          # Skips updating normalizers in the agent, as it's handled in the learner.
          update_normalizers_in_train=False,
          debug_summaries=debug_summaries,
          summarize_grads_and_vars=summarize_grads_and_vars,
          train_step_counter=train_step,
      )
      agent.initialize()
  
      reverb_server = reverb.Server(
          [
              reverb.Table(  # Replay buffer storing experience for training.
                  name="training_table",
                  sampler=reverb.selectors.Fifo(),
                  remover=reverb.selectors.Fifo(),
                  rate_limiter=reverb.rate_limiters.MinSize(1),
                  max_size=replay_capacity,
                  max_times_sampled=1,
              ),
              reverb.Table(  # Replay buffer storing experience for normalization.
                  name="normalization_table",
                  sampler=reverb.selectors.Fifo(),
                  remover=reverb.selectors.Fifo(),
                  rate_limiter=reverb.rate_limiters.MinSize(1),
                  max_size=replay_capacity,
                  max_times_sampled=1,
              ),
          ],
          port=reverb_port,
      )
  
      # Create the replay buffer.
      reverb_replay_train = reverb_replay_buffer.ReverbReplayBuffer(
          agent.collect_data_spec,
          sequence_length=collect_sequence_length,
          table_name="training_table",
          server_address="localhost:{}".format(reverb_server.port),
          # The only collected sequence is used to populate the batches.
          max_cycle_length=1,
          rate_limiter_timeout_ms=1000,
      )
      reverb_replay_normalization = reverb_replay_buffer.ReverbReplayBuffer(
          agent.collect_data_spec,
          sequence_length=collect_sequence_length,
          table_name="normalization_table",
          server_address="localhost:{}".format(reverb_server.port),
          # The only collected sequence is used to populate the batches.
          max_cycle_length=1,
          rate_limiter_timeout_ms=1000,
      )
  
      rb_observer = reverb_utils.ReverbTrajectorySequenceObserver(
          reverb_replay_train.py_client,
          ["training_table", "normalization_table"],
          sequence_length=collect_sequence_length,
          stride_length=collect_sequence_length,
      )
  
      saved_model_dir = os.path.join(
          root_dir, learner.POLICY_SAVED_MODEL_DIR
      )
      collect_env_step_metric = py_metrics.EnvironmentSteps()
      learning_triggers = [
          triggers.PolicySavedModelTrigger(
              saved_model_dir,
              agent,
              train_step,
              interval=policy_save_interval,
              metadata_metrics={
                  triggers.ENV_STEP_METADATA_KEY: collect_env_step_metric
              },
          ),
          triggers.StepPerSecondLogTrigger(
              train_step, interval=summary_interval
          ),
      ]
  
      def training_dataset_fn():
          return reverb_replay_train.as_dataset(
              sample_batch_size=num_environments,
              sequence_preprocess_fn=agent.preprocess_sequence,
          )
  
      def normalization_dataset_fn():
          return reverb_replay_normalization.as_dataset(
              sample_batch_size=num_environments,
              sequence_preprocess_fn=agent.preprocess_sequence,
          )
  
      agent_learner = ppo_learner.PPOLearner(
          root_dir,
          train_step,
          agent,
          experience_dataset_fn=training_dataset_fn,
          normalization_dataset_fn=normalization_dataset_fn,
          num_samples=1,
          num_epochs=num_epochs,
          minibatch_size=minibatch_size,
          shuffle_buffer_size=collect_sequence_length,
          triggers=learning_triggers,
      )
  
      tf_collect_policy = agent.collect_policy
      collect_policy = py_tf_eager_policy.PyTFEagerPolicy(
          tf_collect_policy, use_tf_function=True
      )
  
      collect_actor = actor.Actor(
          collect_env,
          collect_policy,
          train_step,
          steps_per_run=collect_sequence_length,
          observers=[rb_observer],
          metrics=actor.collect_metrics(buffer_size=10)
          + [collect_env_step_metric],
          reference_metrics=[collect_env_step_metric],
          summary_dir=os.path.join(root_dir, learner.TRAIN_DIR),
          summary_interval=summary_interval,
      )
  
      eval_greedy_policy = py_tf_eager_policy.PyTFEagerPolicy(
          agent.policy, use_tf_function=True
      )
  
      if eval_interval:
          logging.info("Intial evaluation.")
          eval_actor = actor.Actor(
              eval_env,
              eval_greedy_policy,
              train_step,
              metrics=actor.eval_metrics(eval_episodes),
              reference_metrics=[collect_env_step_metric],
              summary_dir=os.path.join(root_dir, "eval"),
              episodes_per_run=eval_episodes,
          )
  
          eval_actor.run_and_log()
  
      logging.info("Training on %s", env_name)
      last_eval_step = 0
      for i in range(num_iterations):
          collect_actor.run()
          rb_observer.flush()
          agent_learner.run()
          reverb_replay_train.clear()
          reverb_replay_normalization.clear()
          current_iteration.assign_add(1)
  
          # Eval only if `eval_interval` has been set. Then, eval if the current train
          # step is equal or greater than the `last_eval_step` + `eval_interval` or if
          # this is the last iteration. This logic exists because agent_learner.run()
          # does not return after every train step.
          if eval_interval and (
              agent_learner.train_step_numpy
              >= eval_interval + last_eval_step
              or i == num_iterations - 1
          ):
              logging.info("Evaluating.")
              eval_actor.run_and_log()
              last_eval_step = agent_learner.train_step_numpy
  
      rb_observer.close()
      reverb_server.stop()
  
file_path: agents/tf_agents/experimental/examples/ppo/train_eval_lib.py
imports: import os
  from absl import logging
  import gin
  import reverb
  import tensorflow.compat.v2
  from tf_agents.agents.ppo import ppo_actor_network
  from tf_agents.agents.ppo import ppo_clip_agent
  from tf_agents.environments import suite_mujoco
  from tf_agents.metrics import py_metrics
  from tf_agents.networks import value_network
  from tf_agents.policies import py_tf_eager_policy
  from tf_agents.replay_buffers import reverb_replay_buffer
  from tf_agents.replay_buffers import reverb_utils
  from tf_agents.train import actor
  from tf_agents.train import learner
  from tf_agents.train import ppo_learner
  from tf_agents.train import triggers
  from tf_agents.train.utils import spec_utils
  from tf_agents.train.utils import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: train_eval(root_dir, env_name, num_iterations, actor_fc_layers, value_fc_layers, learning_rate, collect_sequence_length, minibatch_size, num_epochs, importance_ratio_clipping, lambda_value, discount_factor, entropy_regularization, value_pred_loss_coef, use_gae, use_td_lambda_return, gradient_clipping, value_clipping, reverb_port, replay_capacity, policy_save_interval, summary_interval, eval_interval, eval_episodes, debug_summaries, summarize_grads_and_vars) and returns: unknown, reverb_replay_train, reverb_replay_normalization