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()
Check save and restore of policy agent.
trax/rl/training_test.py