added training against checkpoint
This commit is contained in:
parent
ce0588c13b
commit
e815d3f9b2
20 changed files with 86 additions and 68 deletions
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -1,7 +1,7 @@
|
|||
import pytest
|
||||
from generals.agents import ExpanderAgent, HunterAgent, RandomAgent
|
||||
|
||||
from general_bots_training.opponents import SelfPlayOpponent, make_opponent
|
||||
from general_bots_training.opponents import ModelOpponent, SelfPlayOpponent, make_opponent
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -28,3 +28,28 @@ def test_make_opponent_supports_self_play_aliases():
|
|||
def test_make_opponent_rejects_unknown_name():
|
||||
with pytest.raises(ValueError, match="unknown opponent"):
|
||||
make_opponent("turtle")
|
||||
|
||||
|
||||
def test_make_opponent_model_requires_checkpoint():
|
||||
with pytest.raises(ValueError, match="requires a checkpoint path"):
|
||||
make_opponent("model")
|
||||
|
||||
|
||||
def test_make_opponent_model_loads_checkpoint(tmp_path):
|
||||
import equinox as eqx
|
||||
import jax.random as jrandom
|
||||
|
||||
from general_bots_training.network import PolicyValueNetwork
|
||||
|
||||
network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14)
|
||||
checkpoint_path = tmp_path / "model.eqx"
|
||||
eqx.tree_serialise_leaves(checkpoint_path, network)
|
||||
|
||||
opponent = make_opponent(f"model:{checkpoint_path}")
|
||||
|
||||
assert isinstance(opponent, ModelOpponent)
|
||||
|
||||
|
||||
def test_make_opponent_rejects_checkpoint_on_non_model():
|
||||
with pytest.raises(ValueError, match="does not take a checkpoint"):
|
||||
make_opponent("random:some.eqx")
|
||||
|
|
|
|||
Loading…
Reference in a new issue