Improvements to env.reset()

.pre-commit-config.yaml

@@ -21,14 +21,22 @@ repos:
         entry: ruff format
         types: [python]
         require_serial: true
-
+
+      - id: spotlessApply
+        name: format with spotless
+        language: system
+        entry: bash -c 'cd forge && gradle :spotlessApply'
+        pass_filenames: false
+        types: [java]
+
       - id: pyright
         name: type check with pyright
         language: system
         entry: pyright
         types: [python]
         require_serial: true
 
+
 #  # Java hooks
 #  - repo: https://github.com/gherynos/pre-commit-java
 #    rev: v0.6.17

forge/src/main/java/com/mineagent/ClientEventHandler.java

@@ -27,10 +27,7 @@
 import org.lwjgl.opengl.GL11;
 import org.slf4j.Logger;
 
-/**
- * Handles client-side game events and coordinates input injection with
- * observations.
- */
+/** Handles client-side game events and coordinates input injection with observations. */
 public class ClientEventHandler {
   private static final Logger LOGGER = LogUtils.getLogger();
   private static final DataBridge dataBridge = DataBridge.getInstance();
@@ -119,8 +116,7 @@ public static void onClientTick(TickEvent.ClientTickEvent event) {
   }
 
   /**
-   * Handles input suppression when a Python client is connected. Disables the
-   * system cursor to
+   * Handles input suppression when a Python client is connected. Disables the system cursor to
    * prevent real mouse input from interfering.
    */
   private static void handleInputSuppression(Minecraft mc) {
@@ -163,9 +159,7 @@ public static void onPlayerDeath(LivingDeathEvent event) {
     }
   }
 
-  /**
-   * The player the agent controls on this machine (not other players or mobs).
-   */
+  /** The player the agent controls on this machine (not other players or mobs). */
   private static boolean isClientControlledPlayer(LivingEntity entity) {
     return entity instanceof LocalPlayer p && p == Minecraft.getInstance().player;
   }

mineagent/client/connection.py

@@ -16,9 +16,7 @@ class ConnectionConfig:
     action_socket: str = "/tmp/mineagent_action.sock"
     frame_width: int = 320
     frame_height: int = 240
-    timeout: float = 30.0
-    max_retries: int = 3
-    retry_delay: float = 1.0
+    timeout: float = 150.0
 
 
 class AsyncMinecraftClient:
@@ -29,6 +27,8 @@ class AsyncMinecraftClient:
     def __init__(self, config: ConnectionConfig | None = None):
         self.config = config or ConnectionConfig()
         self._observation_reader: asyncio.StreamReader | None = None
+        self._observation_writer: asyncio.StreamWriter | None = None
+        self._action_reader: asyncio.StreamReader | None = None
         self._action_writer: asyncio.StreamWriter | None = None
         self._connected: bool = False
         self._logger = logging.getLogger(__name__)
@@ -37,34 +37,40 @@ def __init__(self, config: ConnectionConfig | None = None):
     def connected(self) -> bool:
         return self._connected
 
-    async def connect(self) -> bool:
-        """Establish connection to the Minecraft Forge mod."""
-        for attempt in range(self.config.max_retries):
+    async def connect(self) -> None:
+        """Open both Unix sockets, waiting for the Forge mod to start them.
+
+        The mod creates the socket files asynchronously with this process, so
+        the first attempts fail with FileNotFoundError (no socket yet) or
+        ConnectionRefusedError (file exists but not listening yet). Retry
+        until self.config.timeout elapses.
+
+        Keep both the reader and writer for each socket. open_unix_connection
+        returns a (reader, writer) pair sharing one transport; discarding
+        either half can let the shared transport be reaped, tearing down the
+        socket while the other half is still in use.
+        """
+        (
+            self._observation_reader,
+            self._observation_writer,
+        ) = await self._open_unix_connection(self.config.observation_socket)
+        self._action_reader, self._action_writer = await self._open_unix_connection(
+            self.config.action_socket
+        )
+        self._connected = True
+
+    async def _open_unix_connection(
+        self, path: str
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        loop = asyncio.get_running_loop()
+        end = loop.time() + self.config.timeout
+        while True:
             try:
-                self._observation_reader, _ = await asyncio.open_unix_connection(
-                    self.config.observation_socket
-                )
-                _, self._action_writer = await asyncio.open_unix_connection(
-                    self.config.action_socket
-                )
-                self._connected = True
-                self._logger.info(
-                    "Connected to Minecraft Forge mod - Observation: %s, Action: %s",
-                    self.config.observation_socket,
-                    self.config.action_socket,
-                )
-                return True
-            except OSError as e:
-                self._logger.warning("Connection attempt %d failed: %s", attempt + 1, e)
-                await self._cleanup()
-                if attempt < self.config.max_retries - 1:
-                    await asyncio.sleep(self.config.retry_delay)
-                else:
-                    self._logger.error(
-                        "Failed to connect after %d attempts", self.config.max_retries
-                    )
-                    return False
-        return False
+                return await asyncio.open_unix_connection(path)
+            except (FileNotFoundError, ConnectionRefusedError):
+                if loop.time() >= end:
+                    raise TimeoutError(f"Timed out waiting for socket {path}")
+                await asyncio.sleep(0.1)
 
     async def disconnect(self) -> None:
         """Disconnect from the Minecraft Forge mod."""
@@ -78,7 +84,12 @@ async def _cleanup(self) -> None:
             self._action_writer.close()
             await self._action_writer.wait_closed()
             self._action_writer = None
+        if self._observation_writer:
+            self._observation_writer.close()
+            await self._observation_writer.wait_closed()
+            self._observation_writer = None
         self._observation_reader = None
+        self._action_reader = None
 
     async def send_action(self, raw_input: RawInput) -> bool:
         """
@@ -128,12 +139,9 @@ async def receive_observation(self) -> Observation:
 
         try:
             header = await self._observation_reader.readexactly(12)
-        except asyncio.IncompleteReadError as e:
+        except asyncio.IncompleteReadError:
             self._connected = False
-            raise ConnectionError(
-                f"Connection lost while reading observation header: "
-                f"got {len(e.partial)} of 12 bytes"
-            ) from e
+            raise
 
         reward = struct.unpack(">d", header[0:8])[0]
         frame_length = struct.unpack(">I", header[8:12])[0]

mineagent/config.py

@@ -20,7 +20,7 @@ class EngineConfig:
         Total number of environment steps before program termination
     """
 
-    image_size: tuple[int, int] = (160, 256)
+    image_size: tuple[int, int] = (240, 320)
     max_steps: int = 10_000
 
 

mineagent/env.py

@@ -1,6 +1,7 @@
 import asyncio
 from dataclasses import dataclass
 from typing import Any
+from pathlib import Path
 
 import gymnasium as gym
 import numpy as np
@@ -61,6 +62,7 @@ def __init__(
 
         self._step_count = 0
         self._last_reward: float = 0.0
+        self._minecraft_process: asyncio.subprocess.Process | None = None
 
     def _ensure_loop(self) -> asyncio.AbstractEventLoop:
         if self._loop is None or self._loop.is_closed():
@@ -71,14 +73,37 @@ def _run_async(self, coro):
         loop = self._ensure_loop()
         return loop.run_until_complete(coro)
 
+    async def _launch_minecraft_process(self) -> None:
+        """Launch minecraft in a Python subprocess.
+
+        Stdout/stderr are discarded because the Forge mod already writes full
+        Log4j logs to ``forge/run/logs/latest.log`` (and ``debug.log``);
+        mirroring them here would just duplicate that.
+        """
+        if (
+            self._minecraft_process is not None
+            and self._minecraft_process.returncode is None
+        ):
+            return  # already running
+
+        self._minecraft_process = await asyncio.create_subprocess_exec(
+            "gradle",
+            "runClient",
+            cwd=Path.cwd() / "forge",
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+            start_new_session=True,
+        )
+
     def reset(
         self, *, seed: int | None = None, options: dict[str, Any] | None = None
     ) -> tuple[np.ndarray, dict[str, Any]]:
         super().reset(seed=seed, options=options)
 
+        self._run_async(self._launch_minecraft_process())
+
         if not self._client.connected:
-            if not self._run_async(self._client.connect()):
-                raise RuntimeError("Failed to connect to Minecraft Forge mod")
+            self._run_async(self._client.connect())
 
         self._run_async(self._client.send_action(RawInput.release_all()))
 
@@ -116,9 +141,24 @@ def step(
     def render(self, mode: str = "rgb_array") -> np.ndarray | None:
         raise NotImplementedError("Rendering is not supported.")
 
+    async def _stop_minecraft_process(self) -> None:
+        if (
+            self._minecraft_process is None
+            or self._minecraft_process.returncode is not None
+        ):
+            return  # already stopped
+
+        self._minecraft_process.terminate()
+        try:
+            await asyncio.wait_for(self._minecraft_process.wait(), timeout=60)
+        except asyncio.TimeoutError:
+            self._minecraft_process.kill()
+            await self._minecraft_process.wait()
+
     def close(self):
         if self._client.connected:
             self._run_async(self._client.disconnect())
+        self._run_async(self._stop_minecraft_process())
         if self._loop and not self._loop.is_closed():
             self._loop.close()
             self._loop = None

mineagent/engine.py → mineagent/run.py

@@ -44,34 +44,36 @@ def run() -> None:
     env = MinecraftEnv(env_config=env_config)
     agent = AgentV1(config.agent)
 
-    frame, info = env.reset()
-    # event_bus.publish(EnvReset(timestamp=datetime.now(), observation=frame))
-    obs = torch.tensor(frame, dtype=torch.float).unsqueeze(0)
-    plt.imshow(frame)
-    plt.show()
-    total_return = 0.0
-    prev_env_reward = 0.0
-    for _ in range(engine_config.max_steps):
-        action = agent.act(obs, reward=prev_env_reward)
-        next_frame, reward, terminated, truncated, info = env.step(action)
-        prev_env_reward = float(reward)
-        next_obs = torch.tensor(next_frame, dtype=torch.float).unsqueeze(0)
-        # event_bus.publish(
-        #     EnvStep(
-        #         timestamp=datetime.now(),
-        #         observation=obs,
-        #         action=action,
-        #         reward=reward,
-        #         next_observation=next_obs,
-        #     )
-        # )
-        total_return += reward
-        obs = next_obs
-        if terminated or truncated:
-            break
+    try:
+        frame, info = env.reset()
+        # event_bus.publish(EnvReset(timestamp=datetime.now(), observation=frame))
+        obs = torch.tensor(frame, dtype=torch.float).unsqueeze(0)
+        plt.imshow(frame)
+        plt.show()
+        total_return = 0.0
+        prev_env_reward = 0.0
+        for _ in range(engine_config.max_steps):
+            action = agent.act(obs, reward=prev_env_reward)
+            next_frame, reward, terminated, truncated, info = env.step(action)
+            prev_env_reward = float(reward)
+            next_obs = torch.tensor(next_frame, dtype=torch.float).unsqueeze(0)
+            # event_bus.publish(
+            #     EnvStep(
+            #         timestamp=datetime.now(),
+            #         observation=obs,
+            #         action=action,
+            #         reward=reward,
+            #         next_observation=next_obs,
+            #     )
+            # )
+            total_return += reward
+            obs = next_obs
+            if terminated or truncated:
+                break
 
-    env.close()
-    # event_bus.publish(Stop(timestamp=datetime.now(), total_return=total_return))
+    finally:
+        env.close()
+        # event_bus.publish(Stop(timestamp=datetime.now(), total_return=total_return))
 
 
 if __name__ == "__main__":

pixi.toml

@@ -28,7 +28,7 @@ matplotlib = "*"
 mineagent = { path = ".", editable = true }
 
 [tasks]
-mineagent = "python -m mineagent.engine"
+mineagent = "python -m mineagent.run"
 gradle-run-client = "cd forge && gradle runClient"
 gradle-build = "cd forge && gradle build"
 gradle-test = "cd forge && gradle test"

tests/client/test_connection.py

@@ -18,8 +18,7 @@ def config():
     return ConnectionConfig(
         frame_width=FRAME_WIDTH,
         frame_height=FRAME_HEIGHT,
-        max_retries=3,
-        retry_delay=0.0,
+        timeout=300,
     )