docstr: Check save and restore of joint AWR trainer.
docstr_tok: ["Check","save","and","restore","of","joint","AWR","trainer","."]
code: def test_awrjoint_save_restore(self):
      """Check save and restore of joint AWR trainer."""
      task = rl_task.RLTask(
          "CartPole-v0", initial_trajectories=2, max_steps=2
      )
      joint_model = functools.partial(
          models.PolicyAndValue,
          body=lambda mode: tl.Serial(tl.Dense(4), tl.Relu()),
      )
      tmp_dir = self.create_tempdir().full_path
      trainer1 = actor_critic_joint.AWRJoint(
          task,
          joint_model=joint_model,
          optimizer=opt.Adam,
          batch_size=4,
          train_steps_per_epoch=1,
          n_trajectories_per_epoch=2,
          output_dir=tmp_dir,
      )
      trainer1.run(2)
      self.assertEqual(trainer1.current_epoch, 2)
      self.assertEqual(trainer1._trainer.step, 2)
      # Agent 2 starts where agent 1 stopped.
      trainer2 = actor_critic_joint.AWRJoint(
          task,
          joint_model=joint_model,
          optimizer=opt.Adam,
          batch_size=4,
          train_steps_per_epoch=1,
          n_trajectories_per_epoch=2,
          output_dir=tmp_dir,
      )
      trainer2.run(1)
      self.assertEqual(trainer2.current_epoch, 3)
      self.assertEqual(trainer2._trainer.step, 3)
      trainer1.close()
      trainer2.close()
  
file_path: trax/trax/rl/actor_critic_joint_test.py
imports: import functools
  from absl.testing import absltest
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax import test_utils
  from trax.rl import actor_critic_joint
  from trax.rl import task
  from trax.supervised import lr_schedules
label_desc: Use image node_id 2 for calling the ActorCriticJointTest obj's underlying member method code with example usage: obj.test_awrjoint_save_restore() without return types