class_docstr: Trains a joint policy-and-value model using actor-critic methods.
class_docstr_tok: ["Trains","a","joint","policy-and-value","model","using","actor-critic","methods","."]
docstr: Chooses an action to play after a trajectory.
docstr_tok: ["Chooses","an","action","to","play","after","a","trajectory","."]
code: 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())
  
imports: import functools
  from trax import data
  from trax import layers
  from trax import supervised
  from trax.fastmath import numpy
  from trax.fastmath import stop_gradient
  from trax.rl import actor_critic
  from trax.rl import distributions
  from trax.rl import rl_layers
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 11 for calling the ActorCriticJointAgent obj's underlying member method code with example usage: obj.policy(trajectory, temperature) and returns: sample,  pred