docstr: Check save and restore of policy agent.
docstr_tok: ["Check","save","and","restore","of","policy","agent","."]
code: def test_policy_gradient_save_restore(self):
      """Check save and restore of policy agent."""
      task = rl_task.RLTask("CartPole-v0", max_steps=2)
      tmp_dir = self.create_tempdir().full_path
      agent1 = training.PolicyGradient(
          task,
          model_fn=self._model_fn,
          optimizer=opt.Adam,
          batch_size=2,
          n_trajectories_per_epoch=2,
          n_eval_episodes=1,
          output_dir=tmp_dir,
      )
      agent1.run(1)
      agent1.run(1)
      self.assertEqual(agent1.current_epoch, 2)
      self.assertEqual(agent1.loop.step, 2)
      # Trainer 2 starts where agent 1 stopped.
      agent2 = training.PolicyGradient(
          task,
          model_fn=self._model_fn,
          optimizer=opt.Adam,
          batch_size=2,
          n_trajectories_per_epoch=2,
          n_eval_episodes=1,
          output_dir=tmp_dir,
      )
      agent2.run(1)
      self.assertEqual(agent2.current_epoch, 3)
      self.assertEqual(agent2.loop.step, 3)
      # Manually set saved epoch to 1.
      dictionary = {
          "epoch": 1,
          "avg_returns": [0.0],
          "avg_returns_temperature_0.0": {200: [0.0]},
      }
      with tf.io.gfile.GFile(
          os.path.join(tmp_dir, "rl.pkl"), "wb"
      ) as f:
          pickle.dump(dictionary, f)
  
      # Trainer 3 restores from a checkpoint with Agent/Loop step mistmatch,
      # should fail.
      def agent3_fn():
          return training.PolicyGradient(
              task,
              model_fn=self._model_fn,
              optimizer=opt.Adam,
              batch_size=2,
              n_trajectories_per_epoch=2,
              n_eval_episodes=1,
              output_dir=tmp_dir,
          )
  
      self.assertRaises(ValueError, agent3_fn)
      agent1.close()
      agent2.close()
  
imports: import functools
  import math
  import os
  import pickle
  from absl.testing import absltest
  import tensorflow
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax import test_utils
  from trax.rl import task
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 4 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_policy_gradient_save_restore() and returns: training