40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
"""
|
|
Unit tests for SearchAgent.
|
|
"""
|
|
|
|
import torch
|
|
import pytest
|
|
from src.search_agent import SearchAgent, PolicyValueNet
|
|
from src.utils import encode_state, get_actions, step
|
|
|
|
|
|
def test_policy_value_net_forward():
|
|
net = PolicyValueNet(input_dim=1, action_space=2)
|
|
x = torch.tensor([[3.0]])
|
|
policy, value = net(x)
|
|
assert policy.shape == (1, 2)
|
|
assert value.shape == (1, 1)
|
|
|
|
|
|
def test_search_agent_action_selection():
|
|
net = PolicyValueNet(input_dim=1, action_space=2)
|
|
agent = SearchAgent(policy_value_net=net, max_depth=2)
|
|
# Start from state 0; actions are 1 and 2
|
|
action = agent.act(0)
|
|
assert action in [1, 2]
|
|
|
|
|
|
def test_search_agent_value_estimation():
|
|
net = PolicyValueNet(input_dim=1, action_space=2)
|
|
agent = SearchAgent(policy_value_net=net, max_depth=3)
|
|
# For state 8, the optimal action is 10 (reward 1)
|
|
action = agent.act(8)
|
|
assert action == 10 or action == 9 # depending on policy, 10 is better
|
|
|
|
|
|
def test_search_agent_terminal_state():
|
|
net = PolicyValueNet(input_dim=1, action_space=2)
|
|
agent = SearchAgent(policy_value_net=net, max_depth=1)
|
|
# State 10 is terminal; agent should return 10
|
|
action = agent.act(10)
|
|
assert action == 10 |