File size: 657 Bytes
5b344d4
 
bafb458
302ae2f
bafb458
 
 
302ae2f
 
 
 
bafb458
 
 
 
 
 
 
 
 
 
 
 
 
 
cf6d475
 
302ae2f
 
 
 
 
 
 
 
 
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
from __future__ import annotations

import enum
from typing import List, Optional

from pydantic import BaseModel

from litrl.env.connect_four import Board

ObservationType = List[Board]


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


class GridResponseType(BaseModel):
    grid: ObservationType
    done: bool


class BotResponseType(GridResponseType):
    action: int