docstr: Check save and restore of A2C trainer.
docstr_tok: ["Check","save","and","restore","of","A2C","trainer","."]
code: def test_a2ctrainer_save_restore(self):
      """Check save and restore of A2C trainer."""
      task = rl_task.RLTask(
          "CartPole-v0", initial_trajectories=0, max_steps=20
      )
      body = lambda mode: tl.Serial(tl.Dense(64), tl.Relu())
      policy_model = functools.partial(models.Policy, body=body)
      value_model = functools.partial(models.Value, body=body)
      tmp_dir = self.create_tempdir().full_path
      trainer1 = actor_critic.A2C(
          task,
          value_model=value_model,
          value_optimizer=opt.Adam,
          value_batch_size=2,
          value_train_steps_per_epoch=1,
          policy_model=policy_model,
          policy_optimizer=opt.Adam,
          policy_batch_size=2,
          policy_train_steps_per_epoch=2,
          n_trajectories_per_epoch=2,
          n_shared_layers=1,
          output_dir=tmp_dir,
      )
      trainer1.run(2)
      self.assertEqual(trainer1.current_epoch, 2)
      self.assertEqual(trainer1._value_trainer.step, 2)
      self.assertEqual(trainer1._policy_trainer.step, 4)
      # Trainer 2 starts where trainer 1 stopped.
      trainer2 = actor_critic.A2C(
          task,
          value_model=value_model,
          value_optimizer=opt.Adam,
          value_batch_size=2,
          value_train_steps_per_epoch=1,
          policy_model=policy_model,
          policy_optimizer=opt.Adam,
          policy_batch_size=2,
          policy_train_steps_per_epoch=2,
          n_trajectories_per_epoch=2,
          n_shared_layers=1,
          output_dir=tmp_dir,
      )
      trainer2.run(1)
      self.assertEqual(trainer2.current_epoch, 3)
      self.assertEqual(trainer2._value_trainer.step, 3)
      self.assertEqual(trainer2._policy_trainer.step, 6)
      trainer1.close()
      trainer2.close()
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax import test_utils
  from trax.rl import actor_critic
  from trax.rl import advantages
  from trax.rl import task
  from trax.supervised import lr_schedules
label_desc: Use image node_id 2 for calling the ActorCriticTest obj's underlying member method code with example usage: obj.test_a2ctrainer_save_restore() without return types