[ENVIRONMENT] Connect4 Environment (#101)

* Configuring the connect4_env using the atari template

* fixing laging behaviour of environment

* Adding rendering logic to the connect 4 environment

* Making changes as requested

* fixing atari example

* addin unittest for connect4_env

---------

Co-authored-by: Davide Testuggine <[email protected]>

Author: GboyeStack-Robotics-ML-Engineer

.github/workflows/docker-build.yml

@@ -79,6 +79,8 @@ jobs:
             dockerfile: src/envs/atari_env/server/Dockerfile
           - name: git-env
             dockerfile: src/envs/git_env/server/Dockerfile
+          - name: my-env  # Add your environment here
+            dockerfile: src/envs/connect4_env/server/Dockerfile
           - name: textarena-env
             dockerfile: src/envs/textarena_env/server/Dockerfile
 

examples/OpenEnv_Tutorial.ipynb

@@ -25,19 +25,24 @@
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 
 from envs.atari_env import AtariEnv, AtariAction
-
+# import envs
+# print(envs.__path__)
 
 def main():
     """Run a simple Atari episode."""
     # Connect to the Atari environment server
     print("Connecting to Atari environment...")
     env = AtariEnv.from_docker_image("ghcr.io/meta-pytorch/openenv-atari-env:latest")
-
+    
+   
     try:
         # Reset the environment
         print("\nResetting environment...")
         result = env.reset()
         print(f"Screen shape: {result.observation.screen_shape}")
+
+    
+    
         print(f"Legal actions: {result.observation.legal_actions}")
         print(f"Lives: {result.observation.lives}")
 

examples/atari_simple.py

@@ -0,0 +1,130 @@
+import sys, os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
+
+
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+from matplotlib.patches import Circle
+from envs.connect4_env import Connect4Action, Connect4Env
+
+
+def render_connect4_board(board, ax, player_colors={1: "red", 2: "yellow", -1: "yellow"}, show=True):
+    """
+    Render a Connect 4 board using matplotlib.
+
+    Args:
+        board: 2D list, numpy array, or board object (6x7) with values:
+               0 -> empty, 1 -> player 1, 2 -> player 2 (or -1 for player 2)
+        player_colors: dict mapping player numbers to colors.
+        show: If True, calls plt.show(). If False, returns the figure.
+
+    Returns:
+        The matplotlib figure and axis (if show=False).
+    """
+    # Extract board data if it's an object with board attribute
+    if hasattr(board, 'board'):
+        b_map = np.array(board.board)
+    elif hasattr(board, '__array__'):
+        b_map = np.array(board)
+    else:
+        b_map = np.array(board)
+    
+    # Handle different player value representations
+    # Some environments use 1 and 2, others use 1 and -1
+    rows, cols = b_map.shape
+
+    ax.set_xlim(0, cols)
+    ax.set_ylim(0, rows)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    # Draw the blue board background
+    rect = plt.Rectangle((0, 0), cols, rows, color="#0055FF", zorder=0)
+    ax.add_patch(rect)
+
+    # Draw circular holes
+    for r in range(rows):
+        for c in range(cols):
+            center = (c + 0.5, rows - 1 - r + 0.5)  # Fixed: removed extra -1
+            val = b_map[r, c]
+            
+            # Handle different value representations
+            if val == 1:
+                color = player_colors[1]
+            elif val == 2 or val == -1:
+                color = player_colors.get(2, player_colors.get(-1, "yellow"))
+            else:
+                color = "white"
+
+            circ = Circle(center, 0.4, color=color, ec="black", lw=1.5)
+            ax.add_patch(circ)
+
+    plt.tight_layout()
+    if show:
+        plt.show()
+    else:
+        return ax
+    
+
+def main(render=True):
+    print("Connecting to Connect4 environment...")
+    env = Connect4Env(base_url="http://localhost:8000")
+
+    try:
+        print("\nResetting environment...")
+        result = env.reset()
+
+        frames = []
+        rewards = []
+        steps = []
+
+        # Collect all frames
+        board = np.array(result.observation.board).reshape(6, 7)
+        frames.append(board.copy())
+        rewards.append(result.reward or 0)
+        steps.append(0)
+
+        for step in range(100):
+            if result.done:
+                break
+
+            action_id = int(np.random.choice(result.observation.legal_actions))
+            result = env.step(Connect4Action(column=action_id))
+
+            board = np.array(result.observation.board).reshape(6, 7)
+            frames.append(board.copy())
+            rewards.append(result.reward or 0)
+            steps.append(step + 1)
+
+            if result.done:
+                print(f"Game finished at step {step + 1} with reward {result.reward}")
+                break
+        
+        if render:
+            # Create a single figure and update it
+            fig, ax = plt.subplots(figsize=(7, 6))
+            
+            def animate_frame(i):
+                ax.clear()
+                # Use the render function but don't show immediately
+                render_connect4_board(frames[i], ax=ax, show=False)
+                ax.set_title(f"Step: {steps[i]}, Reward: {rewards[i]:.2f}\nTotal: {sum(rewards[:i+1]):.2f}", 
+                            fontsize=12, pad=20)
+                return ax.patches
+            
+            # Create animation
+            ani = FuncAnimation(fig, animate_frame, frames=len(frames), 
+                            interval=700, repeat=False, blit=False)
+            
+            plt.tight_layout()
+            plt.show(block=True)
+
+    finally:
+        env.close()
+        print("Environment closed.")
+
+
+if __name__ == "__main__":
+    main(render=True)

examples/connect4.py

@@ -0,0 +1,30 @@
+# 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.
+
+"""
+Connect4 Environment for OpenEnv.
+
+This module provides OpenEnv integration for the classic Connect4 board game.
+
+Example:
+    >>> from envs.Connect4_env import Connect4Env, Connect4Action
+    >>>
+    >>> # Connect to a running server or start via Docker
+    >>> env = Connect4Env.from_docker_image("Connect4-env:latest")
+    >>>
+    >>> # Reset and interact
+    >>> result = env.reset()
+    >>> result = env.step(Connect4Action(column=2)) 
+    >>> print(result.reward, result.done)
+    >>>
+    >>> # Cleanup
+    >>> env.close()
+"""
+
+from .client import Connect4Env
+from .models import Connect4Action, Connect4Observation, Connect4State
+
+__all__ = ["Connect4Env", "Connect4Action", "Connect4Observation", "Connect4State"]

src/envs/connect4_env/README.md

@@ -0,0 +1,99 @@
+# 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.
+
+"""
+Connect4 Environment HTTP Client.
+
+This module provides the client for connecting to a Connect4 Environment server
+over HTTP.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, TYPE_CHECKING
+
+from core.client_types import StepResult
+from core.http_env_client import HTTPEnvClient
+
+from .models import Connect4Action, Connect4Observation, Connect4State
+
+if TYPE_CHECKING:
+    from core.containers.runtime import ContainerProvider
+
+
+class Connect4Env(HTTPEnvClient[Connect4Action, Connect4Observation]):
+    """
+    HTTP client for Connect4 Environment.
+
+    This client connects to a Connect4Environment HTTP server and provides
+    methods to interact with it: reset(), step(), and state access.
+
+    Example:
+        >>> client = Connect4Env(base_url="http://localhost:8000")
+        >>> result = client.reset()
+        >>> print(result.observation.board)
+        >>>
+        >>> # Take an action
+        >>> result = client.step(Connect4Action(column=3))
+        >>> print(result.reward, result.done)
+    """
+
+    def _step_payload(self, action: Connect4Action) -> Dict[str, Any]:
+        """
+        Convert Connect4Action to JSON payload for step request.
+
+        Args:
+            action: Connect4Action instance.
+
+        Returns:
+            Dictionary representation suitable for JSON encoding.
+        """
+        return {
+            "column": action.column,  # column index to drop piece
+        }
+
+    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Connect4Observation]:
+        """
+        Parse server response into StepResult[Connect4Observation].
+
+        Args:
+            payload: JSON response from server.
+
+        Returns:
+            StepResult with Connect4Observation.
+        """
+        obs_data = payload.get("observation", {})
+
+        observation = Connect4Observation(
+            board=obs_data.get("board", [[0]*7 for _ in range(6)]),
+            legal_actions=obs_data.get("legal_actions", []),
+            done=payload.get("done", False),
+            reward=payload.get("reward", 0.0),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward", 0.0),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict[str, Any]) -> Connect4State:
+        """
+        Parse server response into Connect4State object.
+
+        Args:
+            payload: JSON response from /state endpoint.
+
+        Returns:
+            Connect4State object with environment state information.
+        """
+        return Connect4State(
+            episode_id=payload.get("episode_id", ""),
+            board=payload.get("board", [[0]*7 for _ in range(6)]),
+            next_player=payload.get("next_player", 1),
+            step_count=payload.get("step_count", 0),
+        )

src/envs/connect4_env/__init__.py

@@ -0,0 +1,68 @@
+# 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 Connect4 Environment.
+
+This module defines the Action, Observation, and State types for Connect4 games
+via the OpenEnv interface.
+"""
+
+from __future__ import annotations
+from dataclasses import dataclass, field
+import numpy as np
+from typing import List
+
+from core.env_server import Action, Observation, State
+
+
+@dataclass
+class Connect4Action(Action):
+    """
+    Action for Connect4 environment.
+
+    Attributes:
+        column: The column index (0 to 6) where the piece will be placed.
+    """
+    column: int
+
+
+@dataclass(kw_only=True)
+class Connect4Observation(Observation):
+    """
+    Observation for Connect4 environment.
+
+    Attributes:
+        board: The current board as a 2D list (6 rows x 7 columns).
+               1 = current player, -1 = opponent, 0 = empty.
+        legal_actions: List of column indices that are valid moves.
+        done: Whether the game is over.
+        reward: Reward for the last action.
+    """
+    
+    board: List[List[int]]
+    legal_actions: List[int]
+    done: bool = False
+    reward: float = 0.0
+    metadata: dict = field(default_factory=dict)
+    
+
+
+@dataclass(kw_only=True)
+class Connect4State(State):
+    """
+    State for Connect4 environment.
+
+    Attributes:
+        episode_id: Unique ID for the current game.
+        board: Current board state (rows x columns), 0 = empty, 1 = player, -1 = opponent.
+        next_player: Whose turn it is (1 or -1).
+        step_count: Number of steps taken in the game.
+    """
+    episode_id: str
+    board: List[List[int]] = field(default_factory=lambda: np.zeros((6,7), dtype=int).tolist())
+    next_player: int = 1
+    step_count: int = 0

src/envs/connect4_env/client.py

@@ -0,0 +1,18 @@
+ARG BASE_IMAGE=openenv-base:latest
+FROM ${BASE_IMAGE}
+
+# Install any additional dependencies
+RUN pip install --no-cache-dir \
+    gymnasium>=0.29.0 \
+    ale-py>=0.8.0 \
+    numpy>=1.24.0
+# Copy environment code
+COPY src/core/ /app/src/core/
+COPY src/envs/connect4_env/ /app/src/envs/connect4_env/
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run server
+CMD ["uvicorn", "envs.connect4_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]