File size: 3,133 Bytes
d3a24e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """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)
# Should not crash with very low budget
result = planner.plan(
latent=latent,
available_actions=["ACTION1", "ACTION2"],
action_budget_remaining=1,
utility_fn=lambda **kw: 0.0,
)
assert "action" in result
|