What's the best way to deal with variable episode length?

[uduse]

What's the best way to deal with variable episode length?

[ebrevdo]

You can create one table which contains variable lengths. Both the old Writer/Sampler and the new Trajectory Writer/Sampler allow it (if you provide a signature to the dataset, simply set the outer/time dimension to None in Python).

But now you additionally have the problem of batching and training from a stream of variable episode lengths.

You have roughly 3 options from a tf.data perspective:

• padded-batch (what you're using)
• bucket-by-sequence-length (what you're suggesting except move the bucketing into tf.data streams and out of tables; but using multiple tables works too!)
• chunk-and-shuffle (chop all episodes into fixed sized pieces and shuffle those, not bothering to train on full eps)

Chunk-and-shuffle often works well in dense reward scenarios. Here's an example:

  stream = dataset.unbatch()
  # Since we're generally reading from a never-ending replay buffer, we
  # can drop remainder here and get the benefit of usually having known
  # episode lengths and batch sizes.
  chunks = stream.batch(chunk_size, drop_remainder=True)
  # reshuffle_each_iteration will probably never be used since replay
  # buffers are based on never-ending datasets, but we add it just in case.
  shuffled = chunks.shuffle(
      shuffle_buffer_size, reshuffle_each_iteration=True, seed=seed)
  batched = shuffled.batch(batch_size, drop_remainder=True)
  return batched

The bucket-by-sequence-length approach can be combined with padding-up-to-fixed-lengths-per-bucket (allowing fixed shaped graphs, which is useful if you're training with TPUs). The tf.data function that helps you do this is bucket_by_sequence_length.

In all cases, you can accelerate by creating multiple reverb datasets (flexible_batch_size=1 in each) and using tf.data's interleave (ACME's reverb dataset code does this).