def test_sanity_ppo_cartpole(self):
    """Run PPO and check whether it correctly runs for 2 epochs.s."""
    task = rl_task.RLTask(
        'CartPole-v1', initial_trajectories=0, max_steps=200)

    lr = lambda: lr_schedules.multifactor(  # pylint: disable=g-long-lambda
        constant=1e-3,
        warmup_steps=100,
        factors='constant * linear_warmup')

    body = lambda mode: tl.Serial(tl.Dense(64), tl.Relu())
    policy_model = functools.partial(models.Policy, body=body)
    value_model = functools.partial(models.Value, body=body)
    trainer = actor_critic.PPO(
        task,
        n_shared_layers=1,
        value_model=value_model,
        value_optimizer=opt.Adam,
        value_lr_schedule=lr,
        value_batch_size=128,
        value_train_steps_per_epoch=10,
        policy_model=policy_model,
        policy_optimizer=opt.Adam,
        policy_lr_schedule=lr,
        policy_batch_size=128,
        policy_train_steps_per_epoch=10,
        n_trajectories_per_epoch=10)

    trainer.run(2)
    self.assertEqual(2, trainer.current_epoch)
Run PPO and check whether it correctly runs for 2 epochs.s.
trax/rl/actor_critic_test.py