src/typing.py

from __future__ import annotations

import enum
from typing import Optional

from pydantic import BaseModel


class AgentType(enum.Enum):
    RANDOM = "random"
    MCTS = "mcts"
    SAC = "sac"


class RolloutPolicy(enum.Enum):
    RANDOM = "random"
    SAC = "sac"


class CpuConfig(BaseModel):
    agent_type: AgentType
    simulations: Optional[int] = None  # noqa: UP007
    rollout_policy: Optional[RolloutPolicy] = None  # noqa: UP007