def testVTrace(self, batch_size):
    """Tests V-trace against ground truth data calculated in python."""
    seq_len = 5

    values = {
        # Note that this is only for testing purposes using well-formed inputs.
        # In practice we'd be more careful about taking log() of arbitrary
        # quantities.
        'log_rhos':
            np.log((_shaped_arange(seq_len, batch_size)) / batch_size /
                   seq_len + 1),
        # T, B where B_i: [0.9 / (i+1)] * T
        'discounts':
            np.array([[0.9 / (b + 1)
                       for b in range(batch_size)]
                      for _ in range(seq_len)]),
        'rewards':
            _shaped_arange(seq_len, batch_size),
        'values':
            _shaped_arange(seq_len, batch_size) / batch_size,
        'bootstrap_value':
            _shaped_arange(batch_size) + 1.0,
        'clip_rho_threshold':
            3.7,
        'clip_pg_rho_threshold':
            2.2,
    }

    output = vtrace_ops.vtrace_from_importance_weights(**values)

    with self.test_session() as session:
      output_v = session.run(output)

    ground_truth_v = _ground_truth_calculation(**values)
    for a, b in zip(ground_truth_v, output_v):
      self.assertAllClose(a, b)
Tests V-trace against ground truth data calculated in python.
trfl/vtrace_ops_test.py