| """Tests for the world model module.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
|
|
| from agents.wayfinder.world_model import WorldModel, Transition |
|
|
|
|
| class TestWorldModel: |
| """Test cases for WorldModel.""" |
|
|
| @pytest.fixture |
| def model(self) -> WorldModel: |
| """Create a test world model.""" |
| return WorldModel(latent_dim=32, buffer_size=1000, device="cpu") |
|
|
| @pytest.fixture |
| def sample_latent(self) -> np.ndarray: |
| """Create a sample latent vector.""" |
| return np.random.randn(32).astype(np.float32) |
|
|
| def test_predict_change_returns_probability(self, model: WorldModel, sample_latent: np.ndarray) -> None: |
| """Test that predict_change returns a valid probability.""" |
| prob = model.predict_change(sample_latent, {"action": "ACTION1"}) |
| assert 0.0 <= prob <= 1.0 |
|
|
| def test_predict_next_latent_shape(self, model: WorldModel, sample_latent: np.ndarray) -> None: |
| """Test that predict_next_latent returns correct shape.""" |
| next_latent = model.predict_next_latent(sample_latent, {"action": "ACTION1"}) |
| assert next_latent.shape == (32,) |
|
|
| def test_add_transition_increases_buffer(self, model: WorldModel, sample_latent: np.ndarray) -> None: |
| """Test that add_transition adds to the buffer.""" |
| assert model.buffer_size_current == 0 |
| model.add_transition( |
| state_latent=sample_latent, |
| action={"action": "ACTION1"}, |
| next_latent=sample_latent + 0.1, |
| frame_changed=True, |
| ) |
| assert model.buffer_size_current == 1 |
|
|
| def test_add_transition_deduplicates(self, model: WorldModel, sample_latent: np.ndarray) -> None: |
| """Test that identical transitions are deduplicated.""" |
| for _ in range(3): |
| model.add_transition( |
| state_latent=sample_latent, |
| action={"action": "ACTION1"}, |
| next_latent=sample_latent + 0.1, |
| frame_changed=True, |
| ) |
| assert model.buffer_size_current == 1 |
|
|
| def test_train_step_returns_loss(self, model: WorldModel, sample_latent: np.ndarray) -> None: |
| """Test that train_step returns a loss value after enough data.""" |
| |
| for i in range(70): |
| model.add_transition( |
| state_latent=sample_latent + i * 0.01, |
| action={"action": f"ACTION{i % 4 + 1}"}, |
| next_latent=sample_latent + (i + 1) * 0.01, |
| frame_changed=(i % 2 == 0), |
| ) |
| loss = model.train_step(batch_size=32) |
| assert loss >= 0.0 |
|
|
| def test_confidence_starts_low(self, model: WorldModel) -> None: |
| """Test that confidence is 0 before training.""" |
| assert model.confidence() == 0.0 |
|
|
| def test_prediction_error_zero_on_first_step(self, model: WorldModel) -> None: |
| """Test that prediction error is 0 on the first step.""" |
| latent = np.random.randn(32).astype(np.float32) |
| assert model.prediction_error(None, None, latent) == 0.0 |
|
|
