def testValueSequence(self):
    """Indexing value functions by action with a minibatch of sequences."""
    values = [[[1.1, 1.2, 1.3], [1.4, 1.5, 1.6]],
              [[2.1, 2.2, 2.3], [2.4, 2.5, 2.6]],
              [[3.1, 3.2, 3.3], [3.4, 3.5, 3.6]],
              [[4.1, 4.2, 4.3], [4.4, 4.5, 4.6]]]
    action_indices = [[0, 2],
                      [1, 0],
                      [2, 1],
                      [0, 2]]
    result = indexing_ops.batched_index(values, action_indices)
    expected_result = [[1.1, 1.6],
                       [2.2, 2.4],
                       [3.3, 3.5],
                       [4.1, 4.6]]

    with self.test_session() as sess:
      self.assertAllClose(sess.run(result), expected_result)
Indexing value functions by action with a minibatch of sequences.
trfl/indexing_ops_test.py