synth-ai 0.2.4.dev4__py3-none-any.whl → 0.2.4.dev6__py3-none-any.whl

synth_ai/environments/examples/tictactoe/engine.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+import numpy as np
+from dataclasses import dataclass
+from typing import Dict, Any, Optional, Tuple
+
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.environments.reproducibility.core import IReproducibleEngine
+from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.tasks.core import TaskInstance
+
+
+# Action mapping: coordinate strings to board indices
+COORD_TO_IDX = {
+    "A1": 0,
+    "A2": 1,
+    "A3": 2,
+    "B1": 3,
+    "B2": 4,
+    "B3": 5,
+    "C1": 6,
+    "C2": 7,
+    "C3": 8,
+}
+IDX_TO_COORD = {v: k for k, v in COORD_TO_IDX.items()}
+
+# Win condition patterns (row, col, diagonal indices)
+WIN_PATTERNS = [
+    [0, 1, 2],
+    [3, 4, 5],
+    [6, 7, 8],  # rows
+    [0, 3, 6],
+    [1, 4, 7],
+    [2, 5, 8],  # columns
+    [0, 4, 8],
+    [2, 4, 6],  # diagonals
+]
+
+# Player mappings
+PLAYER_MARKS = {"X": 1, "O": 2}
+MARK_TO_PLAYER = {1: "X", 2: "O", 0: " "}
+
+
+@dataclass
+class TicTacToePublicState:
+    board: np.ndarray  # 3x3 array: 0=empty, 1=X, 2=O
+    current_player: str  # "X" or "O"
+    last_move: Optional[str]  # "A1", "B2", etc.
+    winner: Optional[str]  # None, "X", "O", or "draw"
+    move_count: int  # Number of moves made
+    max_moves: int  # Always 9 for TicTacToe
+    terminated: bool  # Game finished
+
+    def diff(self, prev_state: "TicTacToePublicState") -> Dict[str, Any]:
+        differences = {}
+        if not np.array_equal(self.board, prev_state.board):
+            differences["board"] = self.board.tolist()
+        if self.current_player != prev_state.current_player:
+            differences["current_player"] = self.current_player
+        if self.last_move != prev_state.last_move:
+            differences["last_move"] = self.last_move
+        if self.winner != prev_state.winner:
+            differences["winner"] = self.winner
+        if self.move_count != prev_state.move_count:
+            differences["move_count"] = self.move_count
+        if self.terminated != prev_state.terminated:
+            differences["terminated"] = self.terminated
+        return differences
+
+    @property
+    def board_text(self) -> str:
+        lines = []
+        lines.append("  A B C")
+        for i in range(3):
+            row_marks = []
+            for j in range(3):
+                mark = MARK_TO_PLAYER[self.board[i * 3 + j]]
+                row_marks.append(mark)
+            lines.append(f"{i + 1} {' '.join(row_marks)}")
+        return "\n".join(lines)
+
+
+@dataclass
+class TicTacToePrivateState:
+    reward_last: float
+    total_reward: float
+    terminated: bool
+    truncated: bool
+
+    def diff(self, prev_state: "TicTacToePrivateState") -> Dict[str, Any]:
+        differences = {}
+        if self.reward_last != prev_state.reward_last:
+            differences["reward_last"] = self.reward_last
+        if self.total_reward != prev_state.total_reward:
+            differences["total_reward"] = self.total_reward
+        if self.terminated != prev_state.terminated:
+            differences["terminated"] = self.terminated
+        if self.truncated != prev_state.truncated:
+            differences["truncated"] = self.truncated
+        return differences
+
+
+@dataclass
+class TicTacToeEngineSnapshot(StatefulEngineSnapshot):
+    task_instance_dict: Dict
+    engine_snapshot: Dict
+
+
+class TicTacToeWinComponent(RewardComponent):
+    def __init__(self, player_mark: str = "X"):
+        super().__init__()
+        self.player_mark = player_mark
+
+    async def score(self, state: TicTacToePublicState, action: Any) -> float:
+        if state.winner == self.player_mark:
+            return 1.0
+        elif state.winner and state.winner != "draw":
+            return -1.0  # Opponent won
+        return 0.0
+
+
+class TicTacToeDrawComponent(RewardComponent):
+    async def score(self, state: TicTacToePublicState, action: Any) -> float:
+        if state.winner == "draw":
+            return 0.0
+        return 0.0
+
+
+class TicTacToeIllegalMoveComponent(RewardComponent):
+    def __init__(self):
+        self.illegal_move_attempted = False
+
+    async def score(self, state: TicTacToePublicState, action: Any) -> float:
+        if self.illegal_move_attempted:
+            self.illegal_move_attempted = False
+            return -1.0
+        return 0.0
+
+
+class TicTacToeEngine(StatefulEngine, IReproducibleEngine):
+    def __init__(self, task_instance: TaskInstance):
+        self.task_instance = task_instance
+        self.illegal_move_component = TicTacToeIllegalMoveComponent()
+
+        # Determine which player the agent is controlling
+        agent_player = "X"  # Default to X
+        if hasattr(task_instance, "metadata") and hasattr(
+            task_instance.metadata, "starting_player"
+        ):
+            agent_player = task_instance.metadata.starting_player
+
+        self.reward_stack = RewardStack(
+            [
+                TicTacToeWinComponent(player_mark=agent_player),
+                TicTacToeDrawComponent(),
+                self.illegal_move_component,
+            ]
+        )
+
+        # Initialize game state
+        self.board = np.zeros(9, dtype=int)
+        self.current_player = "X"
+        self.last_move = None
+        self.winner = None
+        self.move_count = 0
+        self.terminated = False
+        self.total_reward = 0.0
+
+        # Apply any pre-moves from task instance metadata
+        if hasattr(task_instance, "metadata") and hasattr(task_instance.metadata, "opening_moves"):
+            for move in task_instance.metadata.opening_moves:
+                self._apply_move(move)
+
+    async def _reset_engine(
+        self, *, seed: int | None = None
+    ) -> Tuple[TicTacToePrivateState, TicTacToePublicState]:
+        self.board = np.zeros(9, dtype=int)
+        self.current_player = "X"
+        self.last_move = None
+        self.winner = None
+        self.move_count = 0
+        self.terminated = False
+        self.total_reward = 0.0
+
+        # Apply any pre-moves from task instance metadata
+        if hasattr(self.task_instance, "metadata") and hasattr(
+            self.task_instance.metadata, "opening_moves"
+        ):
+            for move in self.task_instance.metadata.opening_moves:
+                self._apply_move(move)
+
+        public_state = TicTacToePublicState(
+            board=self.board.copy(),
+            current_player=self.current_player,
+            last_move=self.last_move,
+            winner=self.winner,
+            move_count=self.move_count,
+            max_moves=9,
+            terminated=self.terminated,
+        )
+
+        private_state = TicTacToePrivateState(
+            reward_last=0.0,
+            total_reward=self.total_reward,
+            terminated=self.terminated,
+            truncated=False,
+        )
+
+        return private_state, public_state
+
+    async def _step_engine(self, action: str) -> Tuple[TicTacToePrivateState, TicTacToePublicState]:
+        # Validate and apply move
+        if not self._is_valid_move(action, self.board):
+            self.illegal_move_component.illegal_move_attempted = True
+            self.terminated = True
+        else:
+            self._apply_move(action)
+
+        # Create public state
+        public_state = TicTacToePublicState(
+            board=self.board.copy(),
+            current_player=self.current_player,
+            last_move=self.last_move,
+            winner=self.winner,
+            move_count=self.move_count,
+            max_moves=9,
+            terminated=self.terminated,
+        )
+
+        # Calculate rewards
+        reward = await self.reward_stack.step_reward(public_state, action)
+        self.total_reward += reward
+
+        # Create private state
+        private_state = TicTacToePrivateState(
+            reward_last=reward,
+            total_reward=self.total_reward,
+            terminated=self.terminated,
+            truncated=False,
+        )
+
+        return private_state, public_state
+
+    def _apply_move(self, coord: str):
+        if coord not in COORD_TO_IDX:
+            return
+
+        idx = COORD_TO_IDX[coord]
+        if self.board[idx] == 0:
+            self.board[idx] = PLAYER_MARKS[self.current_player]
+            self.last_move = coord
+            self.move_count += 1
+
+            # Check for winner
+            self.winner = self._check_winner(self.board)
+
+            # Check if game is over
+            if self.winner is not None or self.move_count >= 9:
+                self.terminated = True
+            else:
+                # Switch players
+                self.current_player = "O" if self.current_player == "X" else "X"
+
+    def _check_winner(self, board: np.ndarray) -> Optional[str]:
+        # Check all win patterns
+        for pattern in WIN_PATTERNS:
+            values = [board[i] for i in pattern]
+            if values[0] != 0 and values[0] == values[1] == values[2]:
+                return MARK_TO_PLAYER[values[0]]
+
+        # Check for draw
+        if np.all(board != 0):
+            return "draw"
+
+        return None
+
+    def _is_valid_move(self, coord: str, board: np.ndarray) -> bool:
+        if coord not in COORD_TO_IDX:
+            return False
+        idx = COORD_TO_IDX[coord]
+        return board[idx] == 0
+
+    async def _serialize_engine(self) -> TicTacToeEngineSnapshot:
+        return TicTacToeEngineSnapshot(
+            task_instance_dict=await self.task_instance.serialize(),
+            engine_snapshot={
+                "board": self.board.tolist(),
+                "current_player": self.current_player,
+                "last_move": self.last_move,
+                "winner": self.winner,
+                "move_count": self.move_count,
+                "terminated": self.terminated,
+                "total_reward": self.total_reward,
+            },
+        )
+
+    @classmethod
+    async def _deserialize_engine(cls, snapshot: TicTacToeEngineSnapshot) -> "TicTacToeEngine":
+        task_instance = await TaskInstance.deserialize(snapshot.task_instance_dict)
+        engine = cls(task_instance)
+
+        # Restore state
+        engine.board = np.array(snapshot.engine_snapshot["board"])
+        engine.current_player = snapshot.engine_snapshot["current_player"]
+        engine.last_move = snapshot.engine_snapshot["last_move"]
+        engine.winner = snapshot.engine_snapshot["winner"]
+        engine.move_count = snapshot.engine_snapshot["move_count"]
+        engine.terminated = snapshot.engine_snapshot["terminated"]
+        engine.total_reward = snapshot.engine_snapshot["total_reward"]
+
+        return engine
+
+    def get_current_states_for_observation(
+        self,
+    ) -> Tuple[TicTacToePrivateState, TicTacToePublicState]:
+        public_state = TicTacToePublicState(
+            board=self.board.copy(),
+            current_player=self.current_player,
+            last_move=self.last_move,
+            winner=self.winner,
+            move_count=self.move_count,
+            max_moves=9,
+            terminated=self.terminated,
+        )
+
+        private_state = TicTacToePrivateState(
+            reward_last=0.0,
+            total_reward=self.total_reward,
+            terminated=self.terminated,
+            truncated=False,
+        )
+
+        return private_state, public_state
+
+
+class SynthTicTacToeObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: TicTacToePublicState, priv: TicTacToePrivateState
+    ) -> InternalObservation:
+        observation: InternalObservation = {
+            "board_text": pub.board_text,
+            "current_player": pub.current_player,
+            "move_count": pub.move_count,
+            "last_move": pub.last_move,
+            "winner": pub.winner,
+            "terminated": pub.terminated,
+            "reward_last": priv.reward_last,
+            "total_reward": priv.total_reward,
+        }
+        return observation
+
+
+class SynthTicTacToeCheckpointObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: TicTacToePublicState, priv: TicTacToePrivateState
+    ) -> InternalObservation:
+        observation: InternalObservation = {
+            "board_text_final": pub.board_text,
+            "winner_final": pub.winner,
+            "move_count_final": pub.move_count,
+            "total_reward": priv.total_reward,
+            "terminated": pub.terminated,
+        }
+        return observation

synth_ai/environments/examples/tictactoe/environment.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+from typing import Dict, Optional, Any, List, Union
+from pydantic import BaseModel
+
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.environment.tools import (
+    AbstractTool,
+    EnvToolCall,
+    ToolResult,
+)
+from synth_ai.environments.tasks.core import TaskInstance
+
+from .engine import (
+    TicTacToeEngine,
+    TicTacToePublicState,
+    TicTacToePrivateState,
+    TicTacToeEngineSnapshot,
+    SynthTicTacToeObservationCallable,
+    SynthTicTacToeCheckpointObservationCallable,
+)
+
+
+class TicTacToeActionInput(BaseModel):
+    letter: str  # "A", "B", or "C"
+    number: int  # 1, 2, or 3
+
+
+class TicTacToeInteractTool(AbstractTool):
+    name = "interact"
+    description = "Place your mark (X or O) in the specified cell using letter (A, B, C) and number (1, 2, 3) coordinates."
+    call_schema = TicTacToeActionInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: TicTacToeEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            # Parse input - now using separate letter and number parameters
+            letter = call.args.get("letter")
+            number = call.args.get("number")
+
+            if not letter or number is None:
+                return ToolResult(
+                    ok=False, error="Both letter and number parameters are required", payload={}
+                )
+
+            # Validate letter
+            if letter not in ["A", "B", "C"]:
+                return ToolResult(
+                    ok=False, error=f"Letter must be A, B, or C, got '{letter}'", payload={}
+                )
+
+            # Validate number
+            if number not in [1, 2, 3]:
+                return ToolResult(
+                    ok=False, error=f"Number must be 1, 2, or 3, got {number}", payload={}
+                )
+
+            # Convert to coordinate string (e.g., "A1", "B2", etc.)
+            action = f"{letter}{number}"
+
+            # Execute action
+            private_state, public_state = await self.engine._step_engine(action)
+
+            return ToolResult(
+                ok=True,
+                payload={"public_state": public_state, "private_state": private_state},
+            )
+        except Exception as e:
+            return ToolResult(ok=False, error=str(e), payload={})
+
+
+class TicTacToeEnvironment(StatefulEnvironment, ReproducibleEnvironment[TicTacToeEngine]):
+    def __init__(
+        self,
+        task_instance: TaskInstance,
+        custom_step_obs: Optional[GetObservationCallable] = None,
+        custom_ckpt_obs: Optional[GetObservationCallable] = None,
+    ):
+        self.name = "TicTacToe"
+        self.task_instance = task_instance
+        self.custom_step_observation_callable = (
+            custom_step_obs or SynthTicTacToeObservationCallable()
+        )
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or SynthTicTacToeCheckpointObservationCallable()
+        )
+        self.engine = TicTacToeEngine(task_instance)
+        self._interact_tool = TicTacToeInteractTool(self.engine)
+
+    async def initialize(self) -> InternalObservation:
+        # Reset engine and return initial observation
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def step(self, tool_calls) -> InternalObservation:
+        # Validate and normalize tool calls
+        validated_call = self.validate_tool_calls(tool_calls)
+
+        # Execute the interact tool
+        result = await self._interact_tool(validated_call)
+
+        if result.ok:
+            priv = result.payload["private_state"]
+            pub = result.payload["public_state"]
+            return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+        else:
+            # Return error observation
+            priv, pub = self.engine.get_current_states_for_observation()
+            return await self._to_observation(
+                priv,
+                pub,
+                self.custom_step_observation_callable,
+                extra_obs={"error": result.error},
+            )
+
+    async def checkpoint(self) -> InternalObservation:
+        # Return checkpoint observation
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    async def terminate(self) -> InternalObservation:
+        # Mark as terminated and return final observation
+        priv, pub = self.engine.get_current_states_for_observation()
+        pub.terminated = True
+        priv.terminated = True
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    def validate_tool_calls(self, tool_calls) -> EnvToolCall:
+        # Handle various input formats
+        if isinstance(tool_calls, EnvToolCall):
+            validated_call = tool_calls
+        elif isinstance(tool_calls, dict):
+            # Handle dict format
+            if "tool" in tool_calls:
+                validated_call = EnvToolCall(
+                    tool=tool_calls["tool"], args=tool_calls.get("args", {})
+                )
+            elif "name" in tool_calls:
+                # Handle legacy format
+                validated_call = EnvToolCall(
+                    tool=tool_calls["name"], args=tool_calls.get("parameters", {})
+                )
+            elif "function" in tool_calls:
+                # Handle OpenAI function call format
+                validated_call = EnvToolCall(
+                    tool=tool_calls["function"]["name"],
+                    args=tool_calls["function"].get("arguments", {}),
+                )
+            else:
+                # Assume it's just parameters
+                validated_call = EnvToolCall(tool="interact", args=tool_calls)
+        elif isinstance(tool_calls, list):
+            # Take first call from list
+            if len(tool_calls) > 0:
+                validated_call = self.validate_tool_calls(tool_calls[0])
+            else:
+                raise ValueError("Empty tool calls list")
+        else:
+            # Try to convert to dict
+            validated_call = EnvToolCall(tool="interact", args={"action": str(tool_calls)})
+
+        # Validate tool name
+        if validated_call.tool != "interact":
+            raise ValueError(f"Unknown tool: {validated_call.tool}")
+
+        # Convert legacy formats to new letter/number format
+        args = validated_call.args
+        if "position" in args:
+            # Convert numeric position (0-8) to letter/number
+            position = args["position"]
+            if position < 0 or position > 8:
+                raise ValueError(f"Position {position} must be between 0 and 8")
+            letter = ["A", "B", "C"][position // 3]
+            number = (position % 3) + 1
+            args = {"letter": letter, "number": number}
+        elif "action" in args:
+            # Convert coordinate string (e.g., "A1") to letter/number
+            action = args["action"]
+            if len(action) != 2:
+                raise ValueError(f"Action '{action}' must be 2 characters (e.g., 'A1')")
+            letter = action[0].upper()
+            try:
+                number = int(action[1])
+            except ValueError:
+                raise ValueError(f"Action '{action}' must have a numeric second character")
+            args = {"letter": letter, "number": number}
+
+        # Validate final letter/number values
+        if "letter" in args and "number" in args:
+            letter = args["letter"]
+            number = args["number"]
+            if letter not in ["A", "B", "C"]:
+                raise ValueError(f"Letter must be A, B, or C, got '{letter}'")
+            if number not in [1, 2, 3]:
+                raise ValueError(f"Number must be 1, 2, or 3, got {number}")
+
+        return EnvToolCall(tool=validated_call.tool, args=args)
+
+    async def _to_observation(
+        self,
+        priv: TicTacToePrivateState,
+        pub: TicTacToePublicState,
+        obs_cb: Optional[GetObservationCallable],
+        extra_obs: Optional[Dict] = None,
+    ) -> InternalObservation:
+        # Convert states to observation using callback
+        if obs_cb:
+            obs = await obs_cb.get_observation(pub, priv)
+        else:
+            obs: InternalObservation = {}
+
+        if extra_obs and isinstance(obs, dict):
+            obs.update(extra_obs)
+
+        return obs
+
+    async def _serialize_engine(self) -> TicTacToeEngineSnapshot:
+        # Delegate to engine serialization
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snapshot: TicTacToeEngineSnapshot, task_instance: TaskInstance
+    ) -> "TicTacToeEnvironment":
+        # Create new environment instance
+        env = cls(task_instance)
+        # Restore engine from snapshot
+        env.engine = await TicTacToeEngine._deserialize_engine(snapshot)
+        # Update tool reference
+        env._interact_tool = TicTacToeInteractTool(env.engine)
+        return env