fix: retry transient init failures in make_factorio_env

fle/env/gym_env/registry.py

@@ -1,4 +1,5 @@
 import os
+import time
 from dataclasses import dataclass, asdict
 from typing import Any, Dict, List, Optional
 
@@ -94,66 +95,116 @@ def get_environment_spec(self, env_id: str) -> Optional[GymEnvironmentSpec]:
 
 
 def make_factorio_env(spec: GymEnvironmentSpec, run_idx: int) -> FactorioGymEnv:
-    """Factory function to create a Factorio gym environment"""
+    """Factory function to create a Factorio gym environment.
+
+    The instance-creation block (FactorioInstance build + task.setup +
+    FactorioGymEnv wrap) is retried up to ``FLE_INIT_RETRIES`` times
+    (default 3) with backoff. Transient failures like
+    ``"Could not save research state"`` happen when a Factorio container
+    is recycled across samples and its RCON server returns malformed
+    data on the first reset; a short wait + fresh attempt clears it.
+
+    Configuration errors (no containers available, invalid run_idx) are
+    raised without retry — they won't fix themselves.
+    """
     # Create task from the task definition
     task = TaskFactory.create_task(spec.task_config_path)
 
-    # Create Factorio instance
-    try:
-        # Check for external server configuration via environment variables
-        address = os.getenv("FACTORIO_SERVER_ADDRESS")
-        tcp_port = os.getenv("FACTORIO_SERVER_PORT")
-
-        if not address and not tcp_port:
-            try:
-                ips, udp_ports, tcp_ports = get_local_container_ips()
-            except ValueError:
-                raise RuntimeError("No Factorio containers available")
-
-            if len(tcp_ports) == 0:
-                raise RuntimeError("No Factorio containers available")
-
-            # Apply port offset for multiple terminal sessions
-            container_idx = PORT_OFFSET + run_idx
-            if container_idx >= len(tcp_ports):
-                raise RuntimeError(
-                    f"Container index {container_idx} (PORT_OFFSET={PORT_OFFSET} + run_idx={run_idx}) exceeds available containers ({len(tcp_ports)})"
-                )
-
-            address = ips[container_idx]
-            tcp_port = tcp_ports[container_idx]
-
-        common_kwargs = {
-            "address": address,
-            "tcp_port": int(tcp_port),
-            "num_agents": spec.num_agents,
-            "fast": True,
-            "cache_scripts": True,
-            "inventory": {},
-            "all_technologies_researched": False,
-        }
-
-        print(f"Using local Factorio container at {address}:{tcp_port}")
-        if spec.num_agents > 1:
-            instance = run_async_safely(A2AFactorioInstance.create(**common_kwargs))
-        else:
-            instance = FactorioInstance(**common_kwargs)
-
-        # Set initial speed and unpause
-        instance.set_speed_and_unpause(10)
-
-        # Setup the task
-        task.setup(instance)
-
-        # Create and return the gym environment
-        env = FactorioGymEnv(
-            instance=instance, task=task, enable_vision=spec.enable_vision
-        )
+    # Resolve server address (no retry — pure config)
+    address = os.getenv("FACTORIO_SERVER_ADDRESS")
+    tcp_port = os.getenv("FACTORIO_SERVER_PORT")
+
+    if not address and not tcp_port:
+        try:
+            ips, udp_ports, tcp_ports = get_local_container_ips()
+        except ValueError:
+            raise RuntimeError("No Factorio containers available")
 
-        return env
+        if len(tcp_ports) == 0:
+            raise RuntimeError("No Factorio containers available")
+
+        # Apply port offset for multiple terminal sessions
+        container_idx = PORT_OFFSET + run_idx
+        if container_idx >= len(tcp_ports):
+            raise RuntimeError(
+                f"Container index {container_idx} (PORT_OFFSET={PORT_OFFSET} + run_idx={run_idx}) exceeds available containers ({len(tcp_ports)})"
+            )
 
-    except Exception as e:
-        raise RuntimeError(f"Failed to create Factorio environment: {e}")
+        address = ips[container_idx]
+        tcp_port = tcp_ports[container_idx]
+
+    common_kwargs = {
+        "address": address,
+        "tcp_port": int(tcp_port),
+        "num_agents": spec.num_agents,
+        "fast": True,
+        "cache_scripts": True,
+        "inventory": {},
+        "all_technologies_researched": False,
+    }
+
+    print(f"Using local Factorio container at {address}:{tcp_port}")
+
+    # Retry the FactorioInstance build + task.setup block. The "Could
+    # not save research state" failure mode happens when a Factorio
+    # container is recycled across samples and its Lua script registry
+    # holds stale state from the previous instance. ``FactorioInstance``
+    # with ``cache_scripts=True`` (the default) compares checksums and
+    # skips re-uploading scripts that already match — but the Lua VM
+    # state behind those scripts can be corrupt. ``cleanup()`` only
+    # closes the RCON connection, it does not reset the server.
+    #
+    # Fix: on every retry attempt after the first, force
+    # ``cache_scripts=False`` so all scripts are freshly re-uploaded
+    # (re-initializing their global Lua state). Plus longer backoff
+    # so containers have time to settle.
+    max_attempts = int(os.getenv("FLE_INIT_RETRIES", "5"))
+    backoff = [int(x) for x in os.getenv("FLE_INIT_BACKOFF", "2,5,10,20,30").split(",")]
+    last_err: Optional[Exception] = None
+    for attempt in range(1, max_attempts + 1):
+        instance = None
+        try:
+            attempt_kwargs = dict(common_kwargs)
+            if attempt > 1:
+                # Force fresh Lua script upload on retry — clears stale
+                # server-side state that survived ``cleanup()``.
+                attempt_kwargs["cache_scripts"] = False
+
+            if spec.num_agents > 1:
+                instance = run_async_safely(A2AFactorioInstance.create(**attempt_kwargs))
+            else:
+                instance = FactorioInstance(**attempt_kwargs)
+
+            # Set initial speed and unpause
+            instance.set_speed_and_unpause(10)
+
+            # Setup the task
+            task.setup(instance)
+
+            # Create and return the gym environment
+            return FactorioGymEnv(
+                instance=instance, task=task, enable_vision=spec.enable_vision
+            )
+
+        except Exception as e:
+            last_err = e
+            print(
+                f"[make_factorio_env] attempt {attempt}/{max_attempts} failed "
+                f"on container {address}:{tcp_port} "
+                f"(cache_scripts={attempt_kwargs.get('cache_scripts')}): {e!r}"
+            )
+            # Best-effort cleanup of partially-built instance.
+            if instance is not None:
+                try:
+                    instance.cleanup()
+                except Exception:  # noqa: BLE001
+                    pass
+            if attempt < max_attempts:
+                wait = backoff[min(attempt - 1, len(backoff) - 1)]
+                print(f"[make_factorio_env] sleeping {wait}s before retry")
+                time.sleep(wait)
+
+    raise RuntimeError(f"Failed to create Factorio environment: {last_err}")
 
 
 def register_all_environments() -> None: