| """Tests for the planner module.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
|
|
| from agents.wayfinder.planner import Planner, PlannerNode |
| from agents.wayfinder.world_model import WorldModel |
| from agents.wayfinder.intrinsic_reward import IntrinsicReward |
|
|
|
|
| class TestPlannerNode: |
| """Test cases for PlannerNode.""" |
|
|
| def test_ucb1_infinite_for_unvisited(self) -> None: |
| """Test that UCB1 is infinite for unvisited nodes.""" |
| node = PlannerNode(latent=np.zeros(32)) |
| assert node.ucb1() == float("inf") |
|
|
| def test_ucb1_finite_after_visit(self) -> None: |
| """Test that UCB1 is finite after a visit.""" |
| parent = PlannerNode(latent=np.zeros(32)) |
| parent.visits = 10 |
| node = PlannerNode(latent=np.zeros(32), parent=parent) |
| node.visits = 1 |
| node.total_value = 0.5 |
| assert node.ucb1() < float("inf") |
|
|
| def test_mean_value(self) -> None: |
| """Test mean value computation.""" |
| node = PlannerNode(latent=np.zeros(32)) |
| node.visits = 4 |
| node.total_value = 2.0 |
| assert node.mean_value == 0.5 |
|
|
| def test_mean_value_zero_visits(self) -> None: |
| """Test mean value is 0 for unvisited nodes.""" |
| node = PlannerNode(latent=np.zeros(32)) |
| assert node.mean_value == 0.0 |
|
|
|
|
| class TestPlanner: |
| """Test cases for Planner.""" |
|
|
| @pytest.fixture |
| def planner(self) -> Planner: |
| """Create a test planner.""" |
| world_model = WorldModel(latent_dim=32, device="cpu") |
| reward_module = IntrinsicReward() |
| return Planner( |
| world_model=world_model, |
| reward_module=reward_module, |
| max_depth=3, |
| max_simulations=10, |
| ) |
|
|
| def test_plan_returns_action_dict(self, planner: Planner) -> None: |
| """Test that plan returns a valid action dict.""" |
| latent = np.random.randn(32).astype(np.float32) |
| result = planner.plan( |
| latent=latent, |
| available_actions=["ACTION1", "ACTION2", "ACTION3"], |
| action_budget_remaining=100, |
| utility_fn=lambda **kw: 0.0, |
| ) |
| assert "action" in result |
| assert result["action"] in ["ACTION1", "ACTION2", "ACTION3"] |
|
|
| def test_plan_fallback_on_no_actions(self, planner: Planner) -> None: |
| """Test that plan falls back when no valid actions.""" |
| latent = np.random.randn(32).astype(np.float32) |
| result = planner.plan( |
| latent=latent, |
| available_actions=["ACTION1"], |
| action_budget_remaining=100, |
| utility_fn=lambda **kw: 0.0, |
| ) |
| assert "action" in result |
|
|
| def test_plan_scales_with_budget(self, planner: Planner) -> None: |
| """Test that planning scales with remaining budget.""" |
| latent = np.random.randn(32).astype(np.float32) |
| |
| result = planner.plan( |
| latent=latent, |
| available_actions=["ACTION1", "ACTION2"], |
| action_budget_remaining=1, |
| utility_fn=lambda **kw: 0.0, |
| ) |
| assert "action" in result |
|
|
