code: def train_eval(
      root_dir,
      strategy: tf.distribute.Strategy,
      env_name="HalfCheetah-v2",
      # Training params
      initial_collect_steps=10000,
      heavy_priming: bool = False,
      heavy_priming_num_iterations=100000,
      weights_update_per_interaction=1,
      scale_lr_using_rr=False,
      scale_reset_freq=False,
      num_iterations=1000000,
      actor_fc_layers=(256, 256),
      critic_obs_fc_layers=None,
      critic_action_fc_layers=None,
      critic_joint_fc_layers=(256, 256),
      # Agent params
      batch_size=256,
      learning_rate_overwrite=None,
      actor_learning_rate=3e-4,
      critic_learning_rate=3e-4,
      alpha_learning_rate=3e-4,
      gamma=0.99,
      target_update_tau=0.005,
      target_update_period=1,
      reward_scale_factor=0.1,
      # Replay params
      reverb_port=None,
      replay_capacity=1000000,
      # Others
      policy_save_interval=1000000,
      replay_buffer_save_interval=100000,
      eval_interval=10000,
      eval_episodes=30,
      debug_summaries=False,
      summarize_grads_and_vars=False,
      width: float = 1.0,
      train_mode_actor: str = "dense",
      train_mode_value: str = "dense",
      weight_decay: float = 0.0,
      actor_critic_widths_str: str = "",
  ):
      """Trains and evaluates SAC."""
      assert FLAGS.is_mujoco
      assert isinstance(weights_update_per_interaction, int)
      assert weights_update_per_interaction >= 1
      if learning_rate_overwrite:
          actor_learning_rate = learning_rate_overwrite
          critic_learning_rate = learning_rate_overwrite
          alpha_learning_rate = learning_rate_overwrite
      if scale_lr_using_rr:
          actor_learning_rate /= weights_update_per_interaction
          critic_learning_rate /= weights_update_per_interaction
          alpha_learning_rate /= weights_update_per_interaction
      # BEGIN_GOOGLE_INTERNAL
      xm_client = xmanager_api.XManagerApi()
      work_unit = xm_client.get_current_work_unit()
      xm_objective_value_train_reward = (
          work_unit.get_measurement_series(label="train_reward")
      )
      xm_objective_value_reward = work_unit.get_measurement_series(
          label="reward"
      )
      # END_GOOGLE_INTERNAL
  
      if actor_critic_widths_str:
          actor_critic_widths = [
              float(s) for s in actor_critic_widths_str.split("_")
          ]
          width_actor = actor_critic_widths[0]
          width_value = actor_critic_widths[1]
      else:
          width_actor = width
          width_value = width
  
      logging.info("Training SAC on: %s", env_name)
      logging.info("SAC params: train mode actor: %s", train_mode_actor)
      logging.info("SAC params: train mode value: %s", train_mode_value)
      logging.info("SAC params: width: %s", width)
      logging.info(
          "SAC params: actor_critic_widths_str: %s",
          actor_critic_widths_str,
      )
      logging.info("SAC params: width_actor: %s", width_actor)
      logging.info("SAC params: width_value: %s", width_value)
      logging.info("SAC params: weight_decay: %s", weight_decay)
  
      collect_env = suite_mujoco.load(env_name)
      eval_env = suite_mujoco.load(env_name)
  
      (
          _,
          action_tensor_spec,
          time_step_tensor_spec,
      ) = spec_utils.get_tensor_specs(collect_env)
  
      actor_net = create_sequential_actor_network(
          actor_fc_layers=actor_fc_layers,
          action_tensor_spec=action_tensor_spec,
          width=width_actor,
          weight_decay=weight_decay,
      )
  
      critic_net = create_sequential_critic_network(
          obs_fc_layer_units=critic_obs_fc_layers,
          action_fc_layer_units=critic_action_fc_layers,
          joint_fc_layer_units=critic_joint_fc_layers,
          width=width_value,
          weight_decay=weight_decay,
      )
  
      if scale_reset_freq:
          reset_freq_scale = weights_update_per_interaction
      else:
          reset_freq_scale = 1
      with strategy.scope():
          train_step = train_utils.create_train_step()
          agent = RecycledSacAgent(
              time_step_spec=time_step_tensor_spec,
              action_spec=action_tensor_spec,
              actor_network=actor_net,
              critic_network=critic_net,
              actor_optimizer=tf.keras.optimizers.Adam(
                  learning_rate=actor_learning_rate
              ),
              critic_optimizer=tf.keras.optimizers.Adam(
                  learning_rate=critic_learning_rate
              ),
              alpha_optimizer=tf.keras.optimizers.Adam(
                  learning_rate=alpha_learning_rate
              ),
              target_update_tau=target_update_tau,
              target_update_period=target_update_period,
              td_errors_loss_fn=tf.math.squared_difference,
              gamma=gamma,
              reward_scale_factor=reward_scale_factor,
              gradient_clipping=None,
              debug_summaries=debug_summaries,
              summarize_grads_and_vars=summarize_grads_and_vars,
              train_step_counter=train_step,
              reset_freq_scale=reset_freq_scale,
          )
  
          agent.initialize()
          logging.info("agent initialized.")
  
      table_name = "uniform_table"
      table = reverb.Table(
          table_name,
          max_size=replay_capacity,
          sampler=reverb.selectors.Uniform(),
          remover=reverb.selectors.Fifo(),
          rate_limiter=reverb.rate_limiters.MinSize(1),
      )
  
      reverb_checkpoint_dir = os.path.join(
          root_dir,
          learner.TRAIN_DIR,
          learner.REPLAY_BUFFER_CHECKPOINT_DIR,
      )
      reverb_checkpointer = (
          reverb.platform.checkpointers_lib.DefaultCheckpointer(
              path=reverb_checkpoint_dir
          )
      )
      reverb_server = reverb.Server(
          [table], port=reverb_port, checkpointer=reverb_checkpointer
      )
      reverb_replay = reverb_replay_buffer.ReverbReplayBuffer(
          agent.collect_data_spec,
          sequence_length=2,
          table_name=table_name,
          local_server=reverb_server,
      )
      rb_observer = reverb_utils.ReverbAddTrajectoryObserver(
          reverb_replay.py_client,
          table_name,
          sequence_length=2,
          stride_length=1,
      )
  
      def experience_dataset_fn():
          return reverb_replay.as_dataset(
              sample_batch_size=batch_size, num_steps=2
          ).prefetch(50)
  
      saved_model_dir = os.path.join(
          root_dir, learner.POLICY_SAVED_MODEL_DIR
      )
      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: env_step_metric
              },
          ),
          triggers.ReverbCheckpointTrigger(
              train_step,
              interval=replay_buffer_save_interval,
              reverb_client=reverb_replay.py_client,
          ),
          triggers.StepPerSecondLogTrigger(train_step, interval=1000),
      ]
  
      agent_learner = learner.Learner(
          root_dir,
          train_step,
          agent,
          experience_dataset_fn,
          triggers=learning_triggers,
          strategy=strategy,
      )
  
      random_policy = random_py_policy.RandomPyPolicy(
          collect_env.time_step_spec(), collect_env.action_spec()
      )
      initial_collect_actor = actor.Actor(
          collect_env,
          random_policy,
          train_step,
          steps_per_run=initial_collect_steps,
          observers=[rb_observer],
      )
      logging.info("Doing initial collect.")
      initial_collect_actor.run()
  
      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=1,
          metrics=actor.collect_metrics(10),
          summary_dir=os.path.join(root_dir, learner.TRAIN_DIR),
          observers=[rb_observer, env_step_metric],
      )
  
      tf_greedy_policy = greedy_policy.GreedyPolicy(agent.policy)
      eval_greedy_policy = py_tf_eager_policy.PyTFEagerPolicy(
          tf_greedy_policy, use_tf_function=True
      )
  
      eval_actor = actor.Actor(
          eval_env,
          eval_greedy_policy,
          train_step,
          episodes_per_run=eval_episodes,
          metrics=actor.eval_metrics(eval_episodes),
          summary_dir=os.path.join(root_dir, "eval"),
      )
  
      average_returns = []
      if eval_interval:
          logging.info("Evaluating.")
          eval_actor.run_and_log()
          for metric in eval_actor.metrics:
              if isinstance(metric, py_metrics.AverageReturnMetric):
                  average_returns.append(
                      metric._buffer.mean()
                  )  # pylint: disable=protected-access
  
      # heavy priming experiments
      if heavy_priming:
          logging.info("Heavy priming.")
          agent_learner.run(iterations=heavy_priming_num_iterations)
  
      logging.info("Training.")
      for _ in range(num_iterations):
          collect_actor.run()
          agent_learner.run(iterations=weights_update_per_interaction)
  
          if (
              eval_interval
              and agent_learner.train_step_numpy % eval_interval == 0
          ):
              logging.info("Evaluating.")
              eval_actor.run_and_log()
              for metric in eval_actor.metrics:
                  if isinstance(metric, py_metrics.AverageReturnMetric):
                      average_returns.append(
                          metric._buffer.mean()
                      )  # pylint: disable=protected-access
  
      # Log last section of evaluation scores for the final metric.
      idx = int(FLAGS.average_last_fraction * len(average_returns))
      avg_return = np.mean(average_returns[-idx:])
      logging.info(
          "Step %d, Average Return: %f",
          env_step_metric.result(),
          avg_return,
      )
  
      rb_observer.close()
      reverb_server.stop()
  
imports: import functools
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  import numpy
  import reverb
  from rigl.rl.tfagents import sparse_tanh_normal_projection_network
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.agents.sac import sac_agent
  from tf_agents.environments import suite_mujoco
  from tf_agents.keras_layers import inner_reshape
  from tf_agents.metrics import py_metrics
  from tf_agents.networks import nest_map
  from tf_agents.networks import sequential
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import py_tf_eager_policy
  from tf_agents.policies import random_py_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 triggers
  from tf_agents.train.utils import spec_utils
  from tf_agents.train.utils import strategy_utils
  from tf_agents.train.utils import train_utils
  from tf_agents.utils import common
  from tf_agents.utils import object_identity
label_desc: Use image node_id 10 for calling a global function with example usage: train_eval(root_dir, strategy, env_name, initial_collect_steps, heavy_priming, heavy_priming_num_iterations, weights_update_per_interaction, scale_lr_using_rr, scale_reset_freq, num_iterations, actor_fc_layers, critic_obs_fc_layers, critic_action_fc_layers, critic_joint_fc_layers, batch_size, learning_rate_overwrite, actor_learning_rate, critic_learning_rate, alpha_learning_rate, gamma, target_update_tau, target_update_period, reward_scale_factor, reverb_port, replay_capacity, policy_save_interval, replay_buffer_save_interval, eval_interval, eval_episodes, debug_summaries, summarize_grads_and_vars, width, train_mode_actor, train_mode_value, weight_decay, actor_critic_widths_str) and returns: reverb_replay