docstr: Smoke test of PolicyGradient.
docstr_tok: ["Smoke","test","of","PolicyGradient","."]
code: def test_policy_gradient_smoke(self):
      """Smoke test of PolicyGradient."""
      task = rl_task.RLTask("CartPole-v0", max_steps=2)
      tmp_dir = self.create_tempdir().full_path
      agent = 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,
      )
      agent.run(1)
      self.assertEqual(agent.current_epoch, 1)
  
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 2 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_policy_gradient_smoke() without return types