def policy(self, trajectory, temperature=1.0):
    """Chooses an action to play after a trajectory."""
    model = self._eval_model
    model.replicate_weights(self._trainer.model_weights)
    # The two lines below along with the copying
    # before return make the TPU happy
    tr_slice = trajectory.suffix(self._max_slice_length)
    trajectory_np = tr_slice.to_np(timestep_to_np=self.task.timestep_to_np)
    # Add batch dimension to trajectory_np and run the model.
    pred = model(trajectory_np.observation[None, ...])[0]
    # Pick element 0 from the batch (the only one), last (current) timestep.
    pred = pred[0, -1, :]
    sample = self._policy_dist.sample(pred, temperature=temperature)
    return (sample.copy(), pred.copy())
Chooses an action to play after a trajectory.
trax/rl/actor_critic_joint.py