def test_jointa2ctrainer_cartpole_transformer(self):
    """Test-runs joint A2C on cartpole with Transformer."""
    task = rl_task.RLTask('CartPole-v0', initial_trajectories=1,
                          max_steps=2)
    body = lambda mode: models.TransformerDecoder(  # pylint: disable=g-long-lambda
        d_model=4, d_ff=4, n_layers=1, n_heads=1, mode=mode)
    joint_model = functools.partial(models.PolicyAndValue, body=body)
    trainer = actor_critic_joint.A2CJoint(
        task,
        joint_model=joint_model,
        optimizer=opt.RMSProp,
        batch_size=4,
        train_steps_per_epoch=2,
        n_trajectories_per_epoch=2)
    trainer.run(2)
    self.assertEqual(2, trainer.current_epoch)
Test-runs joint A2C on cartpole with Transformer.
trax/rl/actor_critic_joint_test.py