models.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Chess Environment.

The chess environment lets an RL agent play White against a random bot (Black).
"""

from dataclasses import dataclass, field
from typing import List, Optional

from openenv_core.env_server.types import Action, Observation


@dataclass(kw_only=True)
class ChessAction(Action):
    """Action for the Chess environment — a UCI move string (e.g. 'e2e4')."""

    move: str


@dataclass(kw_only=True)
class ChessObservation(Observation):
    """Observation from the Chess environment."""

    board_fen: str
    legal_moves: List[str] = field(default_factory=list)
    white_move: str = ""
    black_move: Optional[str] = None
    material_balance: float = 0.0
    game_status: str = "ongoing"
    captured_pieces: List[str] = field(default_factory=list)