code: def train_eval(
      root_dir,
      env_name="HalfCheetah-v2",
      eval_env_name=None,
      env_load_fn=suite_mujoco.load,
      # The SAC paper reported:
      # Hopper and Cartpole results up to 1000000 iters,
      # Humanoid results up to 10000000 iters,
      # Other mujoco tasks up to 3000000 iters.
      num_iterations=3000000,
      actor_fc_layers=(256, 256),
      critic_obs_fc_layers=None,
      critic_action_fc_layers=None,
      critic_joint_fc_layers=(256, 256),
      # Params for collect
      # Follow https://github.com/haarnoja/sac/blob/master/examples/variants.py
      # HalfCheetah and Ant take 10000 initial collection steps.
      # Other mujoco tasks take 1000.
      # Different choices roughly keep the initial episodes about the same.
      initial_collect_steps=10000,
      collect_steps_per_iteration=1,
      replay_buffer_capacity=1000000,
      # Params for target update
      target_update_tau=0.005,
      target_update_period=1,
      # Params for train
      train_steps_per_iteration=1,
      batch_size=256,
      actor_learning_rate=3e-4,
      critic_learning_rate=3e-4,
      alpha_learning_rate=3e-4,
      td_errors_loss_fn=tf.math.squared_difference,
      gamma=0.99,
      reward_scale_factor=0.1,
      gradient_clipping=None,
      use_tf_functions=True,
      # Params for eval
      num_eval_episodes=30,
      eval_interval=10000,
      # Params for summaries and logging
      train_checkpoint_interval=50000,
      policy_checkpoint_interval=50000,
      rb_checkpoint_interval=50000,
      log_interval=1000,
      summary_interval=1000,
      summaries_flush_secs=10,
      debug_summaries=False,
      summarize_grads_and_vars=False,
      eval_metrics_callback=None,
  ):
      """A simple train and eval for SAC."""
      root_dir = os.path.expanduser(root_dir)
      train_dir = os.path.join(root_dir, "train")
      eval_dir = os.path.join(root_dir, "eval")
  
      train_summary_writer = tf.compat.v2.summary.create_file_writer(
          train_dir, flush_millis=summaries_flush_secs * 1000
      )
      train_summary_writer.set_as_default()
  
      eval_summary_writer = tf.compat.v2.summary.create_file_writer(
          eval_dir, flush_millis=summaries_flush_secs * 1000
      )
      eval_metrics = [
          tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes),
          tf_metrics.AverageEpisodeLengthMetric(
              buffer_size=num_eval_episodes
          ),
      ]
  
      global_step = tf.compat.v1.train.get_or_create_global_step()
      with tf.compat.v2.summary.record_if(
          lambda: tf.math.equal(global_step % summary_interval, 0)
      ):
          tf_env = tf_py_environment.TFPyEnvironment(
              env_load_fn(env_name)
          )
          eval_env_name = eval_env_name or env_name
          eval_tf_env = tf_py_environment.TFPyEnvironment(
              env_load_fn(eval_env_name)
          )
  
          time_step_spec = tf_env.time_step_spec()
          observation_spec = time_step_spec.observation
          action_spec = tf_env.action_spec()
  
          actor_net = actor_distribution_network.ActorDistributionNetwork(
              observation_spec,
              action_spec,
              fc_layer_params=actor_fc_layers,
              continuous_projection_net=tanh_normal_projection_network.TanhNormalProjectionNetwork,
          )
          critic_net = critic_network.CriticNetwork(
              (observation_spec, action_spec),
              observation_fc_layer_params=critic_obs_fc_layers,
              action_fc_layer_params=critic_action_fc_layers,
              joint_fc_layer_params=critic_joint_fc_layers,
              kernel_initializer="glorot_uniform",
              last_kernel_initializer="glorot_uniform",
          )
  
          tf_agent = sac_agent.SacAgent(
              time_step_spec,
              action_spec,
              actor_network=actor_net,
              critic_network=critic_net,
              actor_optimizer=tf.compat.v1.train.AdamOptimizer(
                  learning_rate=actor_learning_rate
              ),
              critic_optimizer=tf.compat.v1.train.AdamOptimizer(
                  learning_rate=critic_learning_rate
              ),
              alpha_optimizer=tf.compat.v1.train.AdamOptimizer(
                  learning_rate=alpha_learning_rate
              ),
              target_update_tau=target_update_tau,
              target_update_period=target_update_period,
              td_errors_loss_fn=td_errors_loss_fn,
              gamma=gamma,
              reward_scale_factor=reward_scale_factor,
              gradient_clipping=gradient_clipping,
              debug_summaries=debug_summaries,
              summarize_grads_and_vars=summarize_grads_and_vars,
              train_step_counter=global_step,
          )
          tf_agent.initialize()
  
          # Make the replay buffer.
          replay_buffer = (
              tf_uniform_replay_buffer.TFUniformReplayBuffer(
                  data_spec=tf_agent.collect_data_spec,
                  batch_size=1,
                  max_length=replay_buffer_capacity,
              )
          )
          replay_observer = [replay_buffer.add_batch]
  
          train_metrics = [
              tf_metrics.NumberOfEpisodes(),
              tf_metrics.EnvironmentSteps(),
              tf_metrics.AverageReturnMetric(
                  buffer_size=num_eval_episodes,
                  batch_size=tf_env.batch_size,
              ),
              tf_metrics.AverageEpisodeLengthMetric(
                  buffer_size=num_eval_episodes,
                  batch_size=tf_env.batch_size,
              ),
          ]
  
          eval_policy = greedy_policy.GreedyPolicy(tf_agent.policy)
          initial_collect_policy = random_tf_policy.RandomTFPolicy(
              tf_env.time_step_spec(), tf_env.action_spec()
          )
          collect_policy = tf_agent.collect_policy
  
          train_checkpointer = common.Checkpointer(
              ckpt_dir=train_dir,
              agent=tf_agent,
              global_step=global_step,
              metrics=metric_utils.MetricsGroup(
                  train_metrics, "train_metrics"
              ),
          )
          policy_checkpointer = common.Checkpointer(
              ckpt_dir=os.path.join(train_dir, "policy"),
              policy=eval_policy,
              global_step=global_step,
          )
          rb_checkpointer = common.Checkpointer(
              ckpt_dir=os.path.join(train_dir, "replay_buffer"),
              max_to_keep=1,
              replay_buffer=replay_buffer,
          )
  
          train_checkpointer.initialize_or_restore()
          rb_checkpointer.initialize_or_restore()
  
          initial_collect_driver = (
              dynamic_step_driver.DynamicStepDriver(
                  tf_env,
                  initial_collect_policy,
                  observers=replay_observer + train_metrics,
                  num_steps=initial_collect_steps,
              )
          )
  
          collect_driver = dynamic_step_driver.DynamicStepDriver(
              tf_env,
              collect_policy,
              observers=replay_observer + train_metrics,
              num_steps=collect_steps_per_iteration,
          )
  
          if use_tf_functions:
              initial_collect_driver.run = common.function(
                  initial_collect_driver.run
              )
              collect_driver.run = common.function(collect_driver.run)
              tf_agent.train = common.function(tf_agent.train)
  
          if replay_buffer.num_frames() == 0:
              # Collect initial replay data.
              logging.info(
                  "Initializing replay buffer by collecting experience for %d steps "
                  "with a random policy.",
                  initial_collect_steps,
              )
              initial_collect_driver.run()
  
          results = metric_utils.eager_compute(
              eval_metrics,
              eval_tf_env,
              eval_policy,
              num_episodes=num_eval_episodes,
              train_step=global_step,
              summary_writer=eval_summary_writer,
              summary_prefix="Metrics",
          )
          if eval_metrics_callback is not None:
              eval_metrics_callback(results, global_step.numpy())
          metric_utils.log_metrics(eval_metrics)
  
          time_step = None
          policy_state = collect_policy.get_initial_state(
              tf_env.batch_size
          )
  
          timed_at_step = global_step.numpy()
          time_acc = 0
  
          # Prepare replay buffer as dataset with invalid transitions filtered.
          def _filter_invalid_transition(trajectories, unused_arg1):
              return ~trajectories.is_boundary()[0]
  
          dataset = (
              replay_buffer.as_dataset(
                  sample_batch_size=batch_size, num_steps=2
              )
              .unbatch()
              .filter(_filter_invalid_transition)
              .batch(batch_size)
              .prefetch(5)
          )
          # Dataset generates trajectories with shape [Bx2x...]
          iterator = iter(dataset)
  
          def train_step():
              experience, _ = next(iterator)
              return tf_agent.train(experience)
  
          if use_tf_functions:
              train_step = common.function(train_step)
  
          global_step_val = global_step.numpy()
          while global_step_val < num_iterations:
              start_time = time.time()
              time_step, policy_state = collect_driver.run(
                  time_step=time_step,
                  policy_state=policy_state,
              )
              for _ in range(train_steps_per_iteration):
                  train_loss = train_step()
              time_acc += time.time() - start_time
  
              global_step_val = global_step.numpy()
  
              if global_step_val % log_interval == 0:
                  logging.info(
                      "step = %d, loss = %f",
                      global_step_val,
                      train_loss.loss,
                  )
                  steps_per_sec = (
                      global_step_val - timed_at_step
                  ) / time_acc
                  logging.info("%.3f steps/sec", steps_per_sec)
                  tf.compat.v2.summary.scalar(
                      name="global_steps_per_sec",
                      data=steps_per_sec,
                      step=global_step,
                  )
                  timed_at_step = global_step_val
                  time_acc = 0
  
              for train_metric in train_metrics:
                  train_metric.tf_summaries(
                      train_step=global_step,
                      step_metrics=train_metrics[:2],
                  )
  
              if global_step_val % eval_interval == 0:
                  results = metric_utils.eager_compute(
                      eval_metrics,
                      eval_tf_env,
                      eval_policy,
                      num_episodes=num_eval_episodes,
                      train_step=global_step,
                      summary_writer=eval_summary_writer,
                      summary_prefix="Metrics",
                  )
                  if eval_metrics_callback is not None:
                      eval_metrics_callback(results, global_step_val)
                  metric_utils.log_metrics(eval_metrics)
  
              if global_step_val % train_checkpoint_interval == 0:
                  train_checkpointer.save(global_step=global_step_val)
  
              if global_step_val % policy_checkpoint_interval == 0:
                  policy_checkpointer.save(global_step=global_step_val)
  
              if global_step_val % rb_checkpoint_interval == 0:
                  rb_checkpointer.save(global_step=global_step_val)
          return train_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import time
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  from six.moves import range
  import tensorflow
  from tf_agents.agents.ddpg import critic_network
  from tf_agents.agents.sac import sac_agent
  from tf_agents.agents.sac import tanh_normal_projection_network
  from tf_agents.drivers import dynamic_step_driver
  from tf_agents.environments import suite_mujoco
  from tf_agents.environments import tf_py_environment
  from tf_agents.eval import metric_utils
  from tf_agents.metrics import tf_metrics
  from tf_agents.networks import actor_distribution_network
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import random_tf_policy
  from tf_agents.replay_buffers import tf_uniform_replay_buffer
  from tf_agents.utils import common
label_desc: Use image node_id 1 for calling a global function with example usage: train_eval(root_dir, env_name, eval_env_name, env_load_fn, num_iterations, actor_fc_layers, critic_obs_fc_layers, critic_action_fc_layers, critic_joint_fc_layers, initial_collect_steps, collect_steps_per_iteration, replay_buffer_capacity, target_update_tau, target_update_period, train_steps_per_iteration, batch_size, actor_learning_rate, critic_learning_rate, alpha_learning_rate, td_errors_loss_fn, gamma, reward_scale_factor, gradient_clipping, use_tf_functions, num_eval_episodes, eval_interval, train_checkpoint_interval, policy_checkpoint_interval, rb_checkpoint_interval, log_interval, summary_interval, summaries_flush_secs, debug_summaries, summarize_grads_and_vars, eval_metrics_callback) and returns: train_loss, unknown, tf_agent