code: def make_trajectory_mask(
      batched_traj: trajectory.Trajectory,
  ) -> types.Tensor:
      """Mask boundary trajectories and those with invalid returns and advantages.
  
      Args:
        batched_traj: Trajectory, doubly-batched [batch_dim, time_dim,...]. It must
          be preprocessed already.
  
      Returns:
        A mask, type tf.float32, that is 0.0 for all between-episode Trajectory
          (batched_traj.step_type is LAST) and 0.0 if the return value is
          unavailable.
      """
      # 1.0 for all valid trajectories. 0.0 where between episodes.
      not_between_episodes = ~batched_traj.is_boundary()
  
      # 1.0 for trajectories with valid return values. 0.0 where return and
      # advantage are both 0. This happens to the last item when the experience gets
      # preprocessed, as insufficient information was available for calculating
      # advantages.
      valid_return_value = ~(
          tf.equal(batched_traj.policy_info["return"], 0)
          & tf.equal(batched_traj.policy_info["advantage"], 0)
      )
  
      return tf.cast(
          not_between_episodes & valid_return_value, tf.float32
      )
  
file_path: agents/tf_agents/agents/ppo/ppo_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Sequence
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import utils
  from tf_agents.specs import distribution_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: make_trajectory_mask(batched_traj) and returns: tf