def _main_test(self, batch_size, lambda_):
    """Tests Retrace against ground truth data calculated in python."""
    seq_len = 5
    # Create log_rhos such that rho will span from near-zero to above the
    # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5),
    # so that rho is in approx [0.08, 12.2).
    log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len)
    log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
    values = {
        "discounts":
            np.array(  # T, B where B_i: [0.9 / (i+1)] * T
                [[0.9 / (b + 1)
                  for b in range(batch_size)]
                 for _ in range(seq_len)]),
        "rewards":
            _shaped_arange(seq_len, batch_size),
        "q_values":
            _shaped_arange(seq_len, batch_size) / batch_size,
        "values":
            _shaped_arange(seq_len, batch_size) / batch_size,
        "bootstrap_value":
            _shaped_arange(batch_size) + 1.0,  # B
        "log_rhos":
            log_rhos
    }
    placeholders = {
        key: tf.placeholder(tf.float32, shape=val.shape)
        for key, val in six.iteritems(values)
    }
    placeholders = {
        k: tf.placeholder(dtype=p.dtype, shape=[None] * len(p.shape))
        for k, p in placeholders.items()
    }

    retrace_returns = continuous_retrace_ops.retrace_from_importance_weights(
        lambda_=lambda_, **placeholders)

    feed_dict = {placeholders[k]: v for k, v in values.items()}
    with self.test_session() as sess:
      retrace_outputvalues = sess.run(retrace_returns, feed_dict=feed_dict)

    ground_truth_data = _ground_truth_calculation(lambda_=lambda_, **values)

    self.assertAllClose(ground_truth_data, retrace_outputvalues.qs)
Tests Retrace against ground truth data calculated in python.
trfl/continuous_retrace_ops_test.py