def test_sanity_awrtrainer_transformer_cartpole(self, thresholds):
    """Test-runs AWR on cartpole with Transformer."""
    task = rl_task.RLTask('CartPole-v0', initial_trajectories=2,
                          max_steps=2)
    body = lambda mode: models.TransformerDecoder(  # pylint: disable=g-long-lambda
        d_model=2, d_ff=2, n_layers=1, n_heads=1, mode=mode)
    policy_model = functools.partial(models.Policy, body=body)
    value_model = functools.partial(models.Value, body=body)
    lr = lambda: lr_schedules.multifactor(  # pylint: disable=g-long-lambda
        constant=1e-2, warmup_steps=100, factors='constant * linear_warmup')
    trainer = actor_critic.AWR(
        task,
        thresholds=thresholds,
        n_shared_layers=0,
        max_slice_length=2,
        added_policy_slice_length=1,
        value_model=value_model,
        value_optimizer=opt.Adam,
        value_lr_schedule=lr,
        value_batch_size=2,
        value_train_steps_per_epoch=2,
        policy_model=policy_model,
        policy_optimizer=opt.Adam,
        policy_lr_schedule=lr,
        policy_batch_size=2,
        policy_train_steps_per_epoch=2,
        n_trajectories_per_epoch=1,
        n_eval_episodes=1)
    trainer.run(2)
    self.assertEqual(2, trainer.current_epoch)
Test-runs AWR on cartpole with Transformer.
trax/rl/actor_critic_test.py