class_docstr: Trains a joint policy-and-value model using AWR.
class_docstr_tok: ["Trains","a","joint","policy-and-value","model","using","AWR","."]
docstr: Joint policy and value loss.
docstr_tok: ["Joint","policy","and","value","loss","."]
code: def joint_loss(self):
      """Joint policy and value loss."""
  
      def f(preds, values, returns, actions, mask):
          advantages = jnp.squeeze(
              returns - stop_gradient(values), axis=-1
          )
          logps = self._policy_dist.log_prob(preds, actions)
          awr_loss = actor_critic.AWRLoss(
              beta=self._beta,
              w_max=self._w_max,
              thresholds=self._thresholds,
          )((logps, advantages, jnp.zeros_like(logps), mask))
          l2_value_loss = (
              jnp.mean((returns - values) ** 2) * self._value_loss_coeff
          )
          return awr_loss + l2_value_loss
  
      return tl.Fn("AWRJointLoss", f)
  
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 3 for calling the AWRJoint obj's underlying member method code with example usage: obj.joint_loss() and returns: tl, unknown