import pytest
import torch
from ding.rl_utils import compute_q_retraces


@pytest.mark.unittest
def test_compute_q_retraces():
    T, B, N = 64, 32, 6
    q_values = torch.randn(T + 1, B, N)
    v_pred = torch.randn(T + 1, B, 1)
    rewards = torch.randn(T, B)
    ratio = torch.rand(T, B, N) * 0.4 + 0.8
    assert ratio.max() <= 1.2 and ratio.min() >= 0.8
    weights = torch.rand(T, B)
    actions = torch.randint(0, N, size=(T, B))
    with torch.no_grad():
        q_retraces = compute_q_retraces(q_values, v_pred, rewards, actions, weights, ratio, gamma=0.99)
    assert q_retraces.shape == (T + 1, B, 1)