166 lines
5.2 KiB
Python
166 lines
5.2 KiB
Python
"""Tests for NN models: CriticNN, ActorNN, compute_q, check_positive_definite."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from hpi.nn_models import ActorNN, CriticNN, check_positive_definite, compute_q
|
|
|
|
|
|
@pytest.fixture
|
|
def critic():
|
|
return CriticNN()
|
|
|
|
|
|
@pytest.fixture
|
|
def actor():
|
|
return ActorNN()
|
|
|
|
|
|
class TestCriticNN:
|
|
"""Test suite for CriticNN."""
|
|
|
|
def test_phi_zero(self, critic):
|
|
"""V(0) = 0."""
|
|
x = torch.zeros(1, 2)
|
|
V = critic(x)
|
|
assert abs(V.item()) < 1e-6
|
|
|
|
def test_phi_dim_single(self, critic):
|
|
"""Single sample: V output shape (,)."""
|
|
x = torch.tensor([[0.5, -0.3]], dtype=torch.float32)
|
|
V = critic(x)
|
|
assert V.shape == (1,)
|
|
assert V.ndim == 1
|
|
|
|
def test_phi_dim_batch(self, critic):
|
|
"""Batch input: V output shape (N,)."""
|
|
x = torch.randn(10, 2)
|
|
V = critic(x)
|
|
assert V.shape == (10,)
|
|
|
|
def test_phi_nonnegative(self, critic):
|
|
"""V(x) >= 0 for all x."""
|
|
x = torch.randn(100, 2)
|
|
V = critic(x)
|
|
assert (V >= -1e-6).all()
|
|
|
|
|
|
class TestActorNN:
|
|
"""Test suite for ActorNN."""
|
|
|
|
def test_psi_zero(self, actor):
|
|
"""u(0) is finite (not NaN, not inf)."""
|
|
x = torch.zeros(1, 2)
|
|
u = actor(x)
|
|
assert torch.isfinite(u).all()
|
|
|
|
def test_psi_dim_single(self, actor):
|
|
"""Single sample: u output shape (,)."""
|
|
x = torch.tensor([[0.5, -0.3]], dtype=torch.float32)
|
|
u = actor(x)
|
|
assert u.shape == (1,)
|
|
assert u.ndim == 1
|
|
|
|
def test_psi_dim_batch(self, actor):
|
|
"""Batch input: u output shape (N,)."""
|
|
x = torch.randn(10, 2)
|
|
u = actor(x)
|
|
assert u.shape == (10,)
|
|
|
|
def test_psi_exact_values(self, actor):
|
|
"""u varies with input (not constant)."""
|
|
x1 = torch.tensor([[0.5, 0.3]], dtype=torch.float32)
|
|
x2 = torch.tensor([[-0.8, 1.2]], dtype=torch.float32)
|
|
u1 = actor(x1).item()
|
|
u2 = actor(x2).item()
|
|
assert abs(u1 - u2) > 1e-8
|
|
|
|
|
|
class TestCriticGrad:
|
|
"""Test suite for CriticNN gradient (replaces evaluate_dphi tests)."""
|
|
|
|
@pytest.fixture
|
|
def critic(self):
|
|
return CriticNN()
|
|
|
|
def test_dphi_dim_single(self, critic):
|
|
"""Single sample: grad_V shape (2,)."""
|
|
x = torch.tensor([[0.5, -0.3]], dtype=torch.float32, requires_grad=True)
|
|
V = critic(x)
|
|
grad_V = torch.autograd.grad(V.sum(), x)[0]
|
|
assert grad_V.shape == (1, 2)
|
|
|
|
def test_dphi_dim_batch(self, critic):
|
|
"""Batch: grad_V shape (N, 2)."""
|
|
x = torch.randn(10, 2, requires_grad=True)
|
|
V = critic(x)
|
|
grad_V = torch.autograd.grad(V.sum(), x)[0]
|
|
assert grad_V.shape == (10, 2)
|
|
|
|
def test_dphi_finite_difference(self, critic):
|
|
"""grad_V matches finite difference approximation."""
|
|
x_np = np.array([[0.5, 0.3]], dtype=np.float32)
|
|
x = torch.tensor(x_np, requires_grad=True)
|
|
V = critic(x)
|
|
grad_V = torch.autograd.grad(V.sum(), x)[0].detach().numpy()[0]
|
|
|
|
eps = 1e-4
|
|
for i in range(2):
|
|
dx = np.zeros(2)
|
|
dx[i] = eps
|
|
x_plus = torch.tensor(x_np + dx.reshape(1, 2), dtype=torch.float32)
|
|
x_minus = torch.tensor(x_np - dx.reshape(1, 2), dtype=torch.float32)
|
|
with torch.no_grad():
|
|
V_plus = critic(x_plus).item()
|
|
V_minus = critic(x_minus).item()
|
|
dphi_fd = (V_plus - V_minus) / (2 * eps)
|
|
np.testing.assert_allclose(grad_V[i], dphi_fd, rtol=1e-3, atol=1e-4)
|
|
|
|
def test_dphi_finite_difference_batch(self, critic):
|
|
"""Batch finite difference check for grad_V."""
|
|
np.random.seed(123)
|
|
x_np = np.random.randn(5, 2).astype(np.float32)
|
|
x = torch.tensor(x_np, requires_grad=True)
|
|
V = critic(x)
|
|
grad_V = torch.autograd.grad(V.sum(), x)[0].detach().numpy()
|
|
|
|
eps = 1e-4
|
|
for j in range(5):
|
|
for i in range(2):
|
|
dx = np.zeros(2)
|
|
dx[i] = eps
|
|
x_plus = x_np.copy()
|
|
x_minus = x_np.copy()
|
|
x_plus[j] += dx
|
|
x_minus[j] -= dx
|
|
with torch.no_grad():
|
|
V_plus = critic(torch.tensor(x_plus)).numpy()[j]
|
|
V_minus = critic(torch.tensor(x_minus)).numpy()[j]
|
|
dphi_fd = (V_plus - V_minus) / (2 * eps)
|
|
np.testing.assert_allclose(grad_V[j, i], dphi_fd, rtol=1e-2, atol=1e-3)
|
|
|
|
|
|
class TestComputeQ:
|
|
"""Test suite for compute_q."""
|
|
|
|
def test_Q_value(self):
|
|
"""Q(x) = x1^2 + x2^2."""
|
|
x = torch.tensor([[0.6, -0.8]], dtype=torch.float32)
|
|
q = compute_q(x)
|
|
expected = 0.6**2 + (-0.8)**2
|
|
assert abs(q.item() - expected) < 1e-6
|
|
|
|
def test_Q_batch_dim(self):
|
|
"""Batch Q output shape (N,)."""
|
|
x = torch.randn(10, 2)
|
|
q = compute_q(x)
|
|
assert q.shape == (10,)
|
|
|
|
def test_Q_batch_values(self):
|
|
"""Batch Q values match manual computation."""
|
|
x = torch.tensor([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=torch.float32)
|
|
q = compute_q(x)
|
|
expected = torch.tensor([5.0, 25.0, 0.0])
|
|
assert torch.allclose(q, expected)
|