docstr: Test-runs joint AWR on cartpole with Transformer.
docstr_tok: ["Test-runs","joint","AWR","on","cartpole","with","Transformer","."]
code: def test_jointawrtrainer_cartpole_transformer(self):
      """Test-runs joint AWR 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.AWRJoint(
          task,
          joint_model=joint_model,
          optimizer=opt.Adam,
          batch_size=4,
          train_steps_per_epoch=2,
          n_trajectories_per_epoch=2,
          max_slice_length=2,
      )
      trainer.run(2)
      self.assertEqual(2, trainer.current_epoch)
  
file_path: trax/trax/rl/actor_critic_joint_test.py
imports: import functools
  from absl.testing import absltest
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax import test_utils
  from trax.rl import actor_critic_joint
  from trax.rl import task
  from trax.supervised import lr_schedules
label_desc: Use image node_id 6 for calling the ActorCriticJointTest obj's underlying member method code with example usage: obj.test_jointawrtrainer_cartpole_transformer() without return types