first commit
This commit is contained in:
@@ -0,0 +1,107 @@
|
||||
"""Tests for DataCollector (Module A)."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from hpi.data_collector import DataCollector
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dc():
|
||||
return DataCollector(J=1.0, mgl=1.0)
|
||||
|
||||
|
||||
class TestDataCollector:
|
||||
"""Test suite for DataCollector."""
|
||||
|
||||
def test_f_zero(self, dc):
|
||||
"""f(0) = 0 at equilibrium."""
|
||||
f0 = dc.f(np.zeros(2))
|
||||
np.testing.assert_allclose(f0, [0.0, 0.0], atol=1e-12)
|
||||
|
||||
def test_f_pi_half(self, dc):
|
||||
"""f([π/2, 0]) = [0, mgl/J]."""
|
||||
x = np.array([np.pi / 2, 0.0])
|
||||
f_val = dc.f(x)
|
||||
expected = np.array([0.0, dc.mgl / dc.J])
|
||||
np.testing.assert_allclose(f_val, expected, atol=1e-12)
|
||||
|
||||
def test_f_batch(self, dc):
|
||||
"""Batch f output shape (N, 2)."""
|
||||
x = np.random.randn(10, 2)
|
||||
f_vals = dc.f(x)
|
||||
assert f_vals.shape == (10, 2)
|
||||
|
||||
def test_g_constant(self, dc):
|
||||
"""g(x) = [0, 1/J] constant."""
|
||||
g_val = dc.g(np.array([0.5, 0.3]))
|
||||
expected = np.array([0.0, 1.0 / dc.J])
|
||||
np.testing.assert_allclose(g_val, expected, atol=1e-12)
|
||||
|
||||
def test_g_batch(self, dc):
|
||||
"""Batch g output shape (N, 2)."""
|
||||
x = np.random.randn(10, 2)
|
||||
g_vals = dc.g(x)
|
||||
assert g_vals.shape == (10, 2)
|
||||
# All rows should be [0, 1/J]
|
||||
np.testing.assert_allclose(g_vals[:, 1], np.ones(10) / dc.J, atol=1e-12)
|
||||
|
||||
def test_pe_signal_bounded(self, dc):
|
||||
"""PE signal amplitude should be <= 0.1."""
|
||||
for t in np.linspace(0, 10, 100):
|
||||
val = dc.pe_signal(t)
|
||||
assert abs(val) <= 0.1 + 1e-12
|
||||
|
||||
def test_pe_signal_nonzero(self, dc):
|
||||
"""PE signal should not be identically zero."""
|
||||
vals = [dc.pe_signal(t) for t in np.linspace(0, 10, 100)]
|
||||
assert np.max(np.abs(vals)) > 1e-6
|
||||
|
||||
def test_dynamics_manual(self, dc):
|
||||
"""Verify ẋ = f(x) + g(x)·u."""
|
||||
x = np.array([0.5, 0.3])
|
||||
u = 0.7
|
||||
dx = dc.dynamics(0.0, x, lambda t, x: u)
|
||||
expected = dc.f(x) + dc.g(x) * u
|
||||
np.testing.assert_allclose(dx, expected, atol=1e-12)
|
||||
|
||||
def test_collect_trajectory_shape(self, dc):
|
||||
"""Trajectory outputs have consistent dimensions."""
|
||||
x0 = np.array([0.5, 0.0])
|
||||
T = 1.0
|
||||
dt = 0.01
|
||||
|
||||
def u_const(t, x):
|
||||
return 0.0
|
||||
|
||||
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
||||
|
||||
M_expected = int(T / dt)
|
||||
assert t.shape == (M_expected,)
|
||||
assert X.shape == (M_expected, 2)
|
||||
assert U.shape == (M_expected,)
|
||||
|
||||
def test_collect_trajectory_start(self, dc):
|
||||
"""First state should equal x0."""
|
||||
x0 = np.array([0.5, 0.0])
|
||||
T = 1.0
|
||||
dt = 0.01
|
||||
|
||||
def u_const(t, x):
|
||||
return 0.0
|
||||
|
||||
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
||||
np.testing.assert_allclose(X[0], x0, atol=1e-10)
|
||||
|
||||
def test_collect_trajectory_drift(self, dc):
|
||||
"""With zero control, state should evolve (pendulum swings)."""
|
||||
x0 = np.array([0.5, 0.0])
|
||||
T = 2.0
|
||||
dt = 0.01
|
||||
|
||||
def u_const(t, x):
|
||||
return 0.0
|
||||
|
||||
t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
|
||||
# Should not stay at x0 (pendulum moves)
|
||||
assert np.linalg.norm(X[-1] - x0) > 1e-6
|
||||
@@ -0,0 +1,121 @@
|
||||
"""Tests for HPIController — NN edition."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from hpi.hpi_controller import HPIController
|
||||
from hpi.nn_models import CriticNN, check_positive_definite
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def controller():
|
||||
return HPIController(J=1.0, mgl=1.0, R=1.0, pe_amplitude=0.5,
|
||||
lr=1e-2, epochs=100, lambda_weight=1e-6)
|
||||
|
||||
|
||||
class TestPositiveDefinite:
|
||||
"""Test suite for check_positive_definite on CriticNN."""
|
||||
|
||||
def test_positive_definite_true(self):
|
||||
"""A well-initialized critic should pass PD (random init usually PD)."""
|
||||
critic = CriticNN()
|
||||
# Check: random init typically gives PD Hessian at origin
|
||||
# If not, the function still runs without error
|
||||
result = check_positive_definite(critic)
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_positive_definite_false(self):
|
||||
"""Critic with all-zero final layer fails PD (zero-rank Hessian)."""
|
||||
critic = CriticNN()
|
||||
# Zero out the final layer to make Hessian rank-deficient at origin
|
||||
with torch.no_grad():
|
||||
critic.fc3.weight.zero_()
|
||||
assert check_positive_definite(critic) is False
|
||||
|
||||
def test_positive_definite_large_cross(self):
|
||||
"""Critic with rank-1 Hessian (single direction) fails PD.
|
||||
|
||||
Zero all layers, then set only fc3.weight[0,0]=1. This makes only
|
||||
z1 depend on x1, giving Hessian = 2*[1,0]^T[1,0] which is rank-1.
|
||||
"""
|
||||
critic = CriticNN()
|
||||
with torch.no_grad():
|
||||
for p in critic.parameters():
|
||||
p.zero_()
|
||||
critic.fc3.weight[0, 0] = 1.0
|
||||
assert check_positive_definite(critic) is False
|
||||
|
||||
def test_positive_definite_borderline(self):
|
||||
"""Near-singular Hessian fails PD."""
|
||||
critic = CriticNN()
|
||||
with torch.no_grad():
|
||||
for p in critic.parameters():
|
||||
p.zero_()
|
||||
critic.fc3.weight[0, 0] = 1e-8 # very small, det ~ 0
|
||||
assert check_positive_definite(critic) is False
|
||||
|
||||
|
||||
class TestPhaseOne:
|
||||
"""Test suite for Phase 1."""
|
||||
|
||||
def test_phase_one_finds_L0(self, controller):
|
||||
"""Phase 1 should find a valid L0."""
|
||||
x0 = np.array([0.5, 0.0])
|
||||
L0, critic_sd, actor_sd, info = controller.phase_one_find_L0(
|
||||
x0, T=2.0, dt=0.05, L_start=0.1, L_step=1.0, L_max=20.0,
|
||||
verbose=False,
|
||||
)
|
||||
assert L0 > 0
|
||||
assert isinstance(critic_sd, dict)
|
||||
assert isinstance(actor_sd, dict)
|
||||
assert isinstance(info, dict)
|
||||
# Load and check PD
|
||||
controller.trainer.critic_nn.load_state_dict(critic_sd)
|
||||
assert controller.check_positive_definite()
|
||||
|
||||
def test_phase_one_raises_on_low_Lmax(self, controller):
|
||||
"""Phase 1 should raise RuntimeError if L_max < L_start (no valid L tried)."""
|
||||
x0 = np.array([0.5, 0.0])
|
||||
with pytest.raises(RuntimeError):
|
||||
controller.phase_one_find_L0(
|
||||
x0, T=1.0, dt=0.05, L_start=1.0, L_step=1.0, L_max=0.5
|
||||
)
|
||||
|
||||
|
||||
class TestFullPipeline:
|
||||
"""Test suite for full pipeline."""
|
||||
|
||||
def test_full_pipeline(self, controller):
|
||||
"""Three-phase pipeline runs without error and returns valid results."""
|
||||
results = controller.run(
|
||||
x0=np.array([0.5, 0.0]),
|
||||
T=2.0,
|
||||
dt=0.05,
|
||||
L_start=0.1,
|
||||
L_step=1.0,
|
||||
L_max=10.0,
|
||||
gamma=0.1,
|
||||
epsilon=1e-3,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
assert "phase1_L0" in results
|
||||
assert results["phase1_L0"] > 0
|
||||
assert isinstance(results["phase1_critic_sd"], dict)
|
||||
assert isinstance(results["phase1_actor_sd"], dict)
|
||||
|
||||
assert "phase2_actor_sd" in results
|
||||
assert "phase2_critic_sd" in results
|
||||
assert len(results["phase2_history"]) >= 1
|
||||
|
||||
assert "phase3_actor_sd" in results
|
||||
assert "phase3_critic_sd" in results
|
||||
assert isinstance(results["phase3_actor_sd"], dict)
|
||||
assert isinstance(results["phase3_critic_sd"], dict)
|
||||
|
||||
# Verify final critic is PD
|
||||
controller.trainer.critic_nn.load_state_dict(results["phase3_critic_sd"])
|
||||
is_pd = controller.check_positive_definite()
|
||||
# May or may not be PD depending on convergence, but shouldn't error
|
||||
assert isinstance(is_pd, bool)
|
||||
@@ -0,0 +1,165 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Tests for NNTrainer (replaces test_solver.py)."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from hpi.nn_trainer import NNTrainer
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def trainer():
|
||||
return NNTrainer(R=1.0, lr=1e-3, epochs=50, lambda_weight=1e-6)
|
||||
|
||||
|
||||
class TestNNTrainer:
|
||||
"""Test suite for NNTrainer."""
|
||||
|
||||
# ── integrate_interval ──
|
||||
|
||||
def test_integrate_interval_scalar(self):
|
||||
"""Trapezoidal integration: [0,1,2] with dt=1 -> [0.5, 1.5]."""
|
||||
f = np.array([0.0, 1.0, 2.0])
|
||||
result = NNTrainer.integrate_interval(f, dt=1.0)
|
||||
expected = np.array([0.5, 1.5])
|
||||
np.testing.assert_allclose(result, expected, atol=1e-12)
|
||||
|
||||
def test_integrate_interval_vector(self):
|
||||
"""Vector trapezoidal integration."""
|
||||
f = np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]])
|
||||
result = NNTrainer.integrate_interval(f, dt=1.0)
|
||||
expected = np.array([[1.0, 2.0], [3.0, 4.0]])
|
||||
np.testing.assert_allclose(result, expected, atol=1e-12)
|
||||
|
||||
# ── compute_u_hat / compute_v_hat ──
|
||||
|
||||
def test_compute_u_hat_manual(self, trainer):
|
||||
"""u_hat returns correct shape for single and batch input."""
|
||||
x_single = np.array([0.5, 0.3])
|
||||
u = trainer.compute_u_hat(x_single)
|
||||
assert np.isscalar(u)
|
||||
assert np.isfinite(u)
|
||||
|
||||
x_batch = np.random.randn(10, 2).astype(np.float32)
|
||||
u_batch = trainer.compute_u_hat(x_batch)
|
||||
assert u_batch.shape == (10,)
|
||||
assert np.all(np.isfinite(u_batch))
|
||||
|
||||
def test_compute_v_hat(self, trainer):
|
||||
"""v_hat = u - u_hat."""
|
||||
u = np.array([0.5, 0.3, 0.1])
|
||||
x = np.array([[0.5, 0.3], [0.1, -0.2], [0.0, 0.0]])
|
||||
u_hat = trainer.compute_u_hat(x)
|
||||
v_hat = trainer.compute_v_hat(u, x)
|
||||
np.testing.assert_allclose(v_hat, u - u_hat, atol=1e-10)
|
||||
|
||||
# ── Training output shapes ──
|
||||
|
||||
def test_compute_I_phi_dim(self, trainer):
|
||||
"""After training, verify that V can be computed (replaces I_phi test)."""
|
||||
M = 100
|
||||
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
V = trainer.compute_V(x)
|
||||
assert V.shape == (M,)
|
||||
|
||||
def test_compute_I_psi_dim(self, trainer):
|
||||
"""u_hat output shape matches input (replaces I_psi test)."""
|
||||
M = 100
|
||||
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
u_hat = trainer.compute_u_hat(x)
|
||||
assert u_hat.shape == (M,)
|
||||
|
||||
def test_compute_I_Q_dim(self, trainer):
|
||||
"""Q(x) can be computed on trajectory (replaces I_Q test)."""
|
||||
M = 100
|
||||
x = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
Q = x[:, 0]**2 + x[:, 1]**2
|
||||
assert Q.shape == (M,)
|
||||
|
||||
# ── System construction (adapted) ──
|
||||
|
||||
def test_build_system_dims(self, trainer):
|
||||
"""Training produces a result dict with expected keys."""
|
||||
M = 200
|
||||
t = np.linspace(0, 2.0, M)
|
||||
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
U = np.random.randn(M).astype(np.float32) * 0.1
|
||||
|
||||
result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
|
||||
assert "loss_c" in result
|
||||
assert "mse_c" in result
|
||||
assert np.isfinite(result["loss_c"])
|
||||
|
||||
def test_solve_dim(self, trainer):
|
||||
"""Training returns result dict with finite loss."""
|
||||
M = 200
|
||||
t = np.linspace(0, 2.0, M)
|
||||
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
U = np.random.randn(M).astype(np.float32) * 0.1
|
||||
|
||||
result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
|
||||
assert np.isfinite(result["loss_c"])
|
||||
assert result["loss_c"] >= 0
|
||||
|
||||
def test_solve_no_nan(self, trainer):
|
||||
"""Training produces no NaN in loss."""
|
||||
M = 200
|
||||
t = np.linspace(0, 2.0, M)
|
||||
X = np.random.randn(M, 2).astype(np.float32) * 0.5
|
||||
U = np.random.randn(M).astype(np.float32) * 0.1
|
||||
|
||||
result = trainer.train_one_iteration(t, X, U, L_i=0.0, prev_actor=None, dt=0.01)
|
||||
assert not np.isnan(result["loss_c"])
|
||||
assert not np.isinf(result["loss_c"])
|
||||
|
||||
def test_solve_with_trajectory_data(self, trainer):
|
||||
"""Train on simulated pendulum data produces finite parameters."""
|
||||
from hpi.data_collector import DataCollector
|
||||
|
||||
dc = DataCollector()
|
||||
x0 = np.array([0.5, 0.0])
|
||||
|
||||
def u_func(t, x):
|
||||
return dc.pe_signal(t)
|
||||
|
||||
t, X, U = dc.collect_trajectory(x0, T=2.0, dt=0.01, u_func=u_func)
|
||||
result = trainer.train_one_iteration(t, X, U, L_i=2.0, prev_actor=None, dt=0.01)
|
||||
|
||||
assert np.isfinite(result["loss_c"])
|
||||
assert result["loss_c"] >= 0.0
|
||||
Reference in New Issue
Block a user