def policy(self, trajectory, temperature=1.0):
    """Policy function that allows to play using this agent."""
    tr_slice = trajectory.suffix(self._max_slice_length)
    trajectory_np = tr_slice.to_np(timestep_to_np=self.task.timestep_to_np)
    return rl_training.network_policy(
        collect_model=self._collect_model,
        policy_distribution=self._policy_dist,
        loop=self.loop,
        trajectory_np=trajectory_np,
        head_index=0,
        temperature=temperature,
    )
Policy function that allows to play using this agent.
trax/rl/actor_critic.py