docstr: Trains a policy on cartpole.
docstr_tok: ["Trains","a","policy","on","cartpole","."]
code: def test_policy_gradient_cartpole(self):
      """Trains a policy on cartpole."""
      task = rl_task.RLTask("CartPole-v0", max_steps=200)
      lr = lambda: lr_schedules.multifactor(
          constant=1e-2, factors="constant"
      )
      max_avg_returns = -math.inf
      for _ in range(2):
          agent = training.PolicyGradient(
              task,
              model_fn=self._model_fn,
              optimizer=opt.Adam,
              lr_schedule=lr,
              batch_size=128,
              eval_temperatures=[0.0, 0.5],
              n_eval_episodes=1,
              n_trajectories_per_epoch=2,
          )
          # Assert that we get to 200 at some point and then exit so the test is as
          # fast as possible.
          for ep in range(200):
              agent.run(1)
              self.assertEqual(agent.current_epoch, ep + 1)
              if agent.avg_returns[-1] == 200.0:
                  for eval_t in agent._eval_temperatures:
                      self.assertEqual(
                          len(
                              agent._avg_returns_temperatures[eval_t][
                                  200
                              ]
                          ),
                          len(agent.avg_returns),
                      )
                  return
          max_avg_returns = max(max_avg_returns, agent.avg_returns[-1])
      self.fail(
          "The expected score of 200 has not been reached. "
          "Maximum at end was {}.".format(max_avg_returns)
      )
  
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 5 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_policy_gradient_cartpole() without return types