ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


Create a mask for transitions and optionally final incomplete episodes.

batched_next_time_step Next timestep, doubly-batched [batch_dim, time_dim, ...].
allow_partial_episodes If true, then steps on incomplete episodes are allowed.

A mask, type tf.float32, that is 0.0 for all between-episode timesteps (batched_next_time_step is FIRST). If allow_partial_episodes is set to False, the mask has 0.0 for incomplete episode at the end of the sequence.