def test_jointppotrainer_cartpole(self):
    """Test-runs joint PPO on CartPole."""

    task = rl_task.RLTask('CartPole-v0', initial_trajectories=0,
                          max_steps=2)
    joint_model = functools.partial(
        models.PolicyAndValue,
        body=lambda mode: tl.Serial(tl.Dense(2), tl.Relu()),
    )
    lr = lambda: lr_schedules.multifactor(  # pylint: disable=g-long-lambda
        constant=1e-2, warmup_steps=100, factors='constant * linear_warmup')

    trainer = actor_critic_joint.PPOJoint(
        task,
        joint_model=joint_model,
        optimizer=opt.Adam,
        lr_schedule=lr,
        batch_size=4,
        train_steps_per_epoch=2,
        n_trajectories_per_epoch=5)
    trainer.run(2)
    self.assertEqual(2, trainer.current_epoch)
Test-runs joint PPO on CartPole.
trax/rl/actor_critic_joint_test.py