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

synth_ai/environments/examples/crafter_classic/engine.py

@@ -0,0 +1,579 @@
+"""CrafterEngine — Stateful, reproducible wrapper around danijar/crafter.Env.
+This file follows the same structure as the SokobanEngine shown earlier.
+"""
+
+from __future__ import annotations
+
+# Import logging configuration first to suppress JAX debug messages
+from .config_logging import safe_compare
+
+# Import patches
+from . import engine_deterministic_patch  # Ensures deterministic behavior
+from . import engine_serialization_patch_v3 as engine_serialization_patch  # Adds save/load methods
+from . import world_config_patch_simple as world_config_patch  # Adds configurable world generation
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import crafter  # type: ignore
+import copy
+import dataclasses
+
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.environments.tasks.core import TaskInstance
+from synth_ai.environments.reproducibility.core import IReproducibleEngine
+from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent  # Added
+
+# Local helper imports (must exist relative to this file)
+from .engine_helpers.action_map import CRAFTER_ACTION_MAP  # action‑name → int
+from .engine_helpers.serialization import (
+    serialize_world_object,
+)
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# -----------------------------------------------------------------------------
+# Dataclasses for snapshot & (public, private) runtime state
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class CrafterEngineSnapshot(StatefulEngineSnapshot):
+    env_raw_state: Any  # from crafter.Env.save()
+    total_reward_snapshot: float
+    crafter_seed: Optional[int] = None
+    # Store previous states needed for reward calculation to resume correctly
+    previous_public_state_snapshot: Optional[Dict] = None
+    previous_private_state_snapshot: Optional[Dict] = None
+    # Add _previous_public_state_for_reward and _previous_private_state_for_reward if needed for perfect resume
+    # For RewardStack, its configuration is fixed at init. If it had internal state, that would need saving.
+
+
+@dataclass
+class CrafterPublicState:
+    inventory: Dict[str, int]
+    achievements_status: Dict[str, bool]
+    player_position: Tuple[int, int]
+    player_direction: Union[int, Tuple[int, int]]
+    semantic_map: Optional[np.ndarray]
+    world_material_map: np.ndarray
+    observation_image: np.ndarray
+    num_steps_taken: int
+    max_steps_episode: int
+    error_info: Optional[str] = None
+
+    def diff(self, prev_state: "CrafterPublicState") -> Dict[str, Any]:
+        changes = {}
+        for field in self.__dataclass_fields__:  # type: ignore[attr-defined]
+            new_v, old_v = getattr(self, field), getattr(prev_state, field)
+            if isinstance(new_v, np.ndarray):
+                if not np.array_equal(new_v, old_v):
+                    changes[field] = True
+            elif new_v != old_v:
+                changes[field] = (old_v, new_v)
+        return changes
+
+
+@dataclass
+class CrafterPrivateState:
+    reward_last_step: float
+    total_reward_episode: float
+    achievements_current_values: Dict[str, int]
+    terminated: bool
+    truncated: bool
+    player_internal_stats: Dict[str, Any]
+    world_rng_state_snapshot: Any
+
+    def diff(self, prev_state: "CrafterPrivateState") -> Dict[str, Any]:
+        changes = {}
+        for field in self.__dataclass_fields__:  # type: ignore[attr-defined]
+            new_v, old_v = getattr(self, field), getattr(prev_state, field)
+            if new_v != old_v:
+                changes[field] = (old_v, new_v)
+        return changes
+
+
+# -----------------------------------------------------------------------------
+# Observation helpers
+# -----------------------------------------------------------------------------
+
+
+class CrafterObservationCallable(GetObservationCallable):
+    def __init__(self) -> None:
+        pass
+
+    async def get_observation(
+        self, pub: CrafterPublicState, priv: CrafterPrivateState
+    ) -> InternalObservation:  # type: ignore[override]
+        observation: Dict[str, Any] = {
+            "inventory": pub.inventory,
+            "achievements": pub.achievements_status,
+            "player_pos": pub.player_position,
+            "steps": pub.num_steps_taken,
+            "reward_last": priv.reward_last_step,
+            "total_reward": priv.total_reward_episode,
+            "terminated": priv.terminated,
+            "truncated": priv.truncated,
+        }
+        return observation  # type: ignore[return-value]
+
+
+# -----------------------------------------------------------------------------
+# CrafterEngine implementation
+# -----------------------------------------------------------------------------
+
+
+class CrafterEngine(StatefulEngine, IReproducibleEngine):
+    """StatefulEngine wrapper around `crafter.Env` supporting full snapshotting."""
+
+    task_instance: TaskInstance
+    env: crafter.Env
+
+    # ────────────────────────────────────────────────────────────────────────
+    # Construction helpers
+    # ────────────────────────────────────────────────────────────────────────
+
+    def __init__(self, task_instance: TaskInstance):
+        self.task_instance = task_instance
+        self._total_reward: float = 0.0
+        self._current_action_for_reward: Optional[int] = None
+        self._previous_public_state_for_reward: Optional[CrafterPublicState] = None
+        self._previous_private_state_for_reward: Optional[CrafterPrivateState] = (
+            None  # For stat changes
+        )
+
+        # Initialize achievements tracking
+        self.achievements_unlocked: set = set()
+
+        cfg = getattr(task_instance, "config", {}) or {}
+        area: Tuple[int, int] = tuple(cfg.get("area", (64, 64)))  # type: ignore[arg-type]
+        length: int = int(cfg.get("length", 10000))
+
+        # Get seed from metadata if available, otherwise fall back to config
+        seed: Optional[int] = cfg.get("seed")
+        if hasattr(task_instance, "metadata") and hasattr(task_instance.metadata, "seed"):
+            seed = task_instance.metadata.seed
+
+        # Get world configuration from metadata or config
+        world_config = "normal"  # default
+        world_config_path = None
+
+        if hasattr(task_instance, "metadata") and hasattr(task_instance.metadata, "world_config"):
+            world_config = task_instance.metadata.world_config
+            logger.info(f"CrafterEngine: Using world_config from metadata: {world_config}")
+        elif cfg.get("world_config"):
+            world_config = cfg.get("world_config")
+            logger.info(f"CrafterEngine: Using world_config from cfg: {world_config}")
+
+        if hasattr(task_instance, "metadata") and hasattr(
+            task_instance.metadata, "world_config_path"
+        ):
+            world_config_path = task_instance.metadata.world_config_path
+        elif cfg.get("world_config_path"):
+            world_config_path = cfg.get("world_config_path")
+
+        logger.info(f"CrafterEngine: Creating Env with world_config={world_config}, seed={seed}")
+        self.env = crafter.Env(
+            area=area,
+            length=length,
+            seed=seed,
+            world_config=world_config,
+            world_config_path=world_config_path,
+        )
+        # store original seed for reproducibility
+        self.env._seed = seed
+
+        self.reward_stack = RewardStack(
+            components=[
+                CrafterAchievementComponent(),
+                CrafterPlayerStatComponent(),
+                CrafterStepPenaltyComponent(penalty=-0.001),
+            ]
+        )
+
+    # ────────────────────────────────────────────────────────────────────────
+    # Utility: action validation / mapping
+    # ────────────────────────────────────────────────────────────────────────
+
+    def _validate_action_engine(self, action: Union[int, str]) -> int:  # type: ignore[override]
+        if isinstance(action, str):
+            action = CRAFTER_ACTION_MAP.get(action, 0)
+        if not isinstance(action, int):
+            return 0
+        return int(np.clip(action, 0, len(crafter.constants.actions) - 1))  # type: ignore
+
+    # ────────────────────────────────────────────────────────────────────────
+    # Core StatefulEngine API
+    # ────────────────────────────────────────────────────────────────────────
+
+    async def _reset_engine(
+        self, *, seed: Optional[int] | None = None
+    ) -> Tuple[CrafterPrivateState, CrafterPublicState]:
+        if seed is not None:
+            # Re‑instantiate env with new seed to match crafter's internal reseeding convention
+            self.env = crafter.Env(area=self.env._area, length=self.env._length, seed=seed)
+        obs_img = self.env.reset()
+        self._total_reward = 0.0
+        pub = self._build_public_state(obs_img)
+        priv = self._build_private_state(reward=0.0, terminated=False, truncated=False)
+
+        # Player starting position tracked internally
+
+        return priv, pub
+
+    async def _step_engine(self, action: int) -> Tuple[CrafterPrivateState, CrafterPublicState]:
+        step_start_time = time.time()
+        try:
+            # Validate action is in valid range
+            if action < 0 or action >= self.env.action_space.n:
+                raise ValueError(
+                    f"Invalid action {action}, must be in range [0, {self.env.action_space.n})"
+                )
+
+            # Ensure player reference is valid before proceeding
+            if self.env._player is None:
+                # Try to find player in world objects
+                for obj in self.env._world._objects:
+                    if (
+                        obj is not None
+                        and hasattr(obj, "__class__")
+                        and obj.__class__.__name__ == "Player"
+                    ):
+                        self.env._player = obj
+                        break
+
+                if self.env._player is None:
+                    raise RuntimeError("Player object not found in world")
+
+            # Build public state BEFORE step (baseline if needed)
+            pub_state_before = self._build_public_state(self.env.render())
+
+            # Step the environment
+            crafter_step_start = time.time()
+            obs, reward, done, info = self.env.step(action)
+            crafter_step_time = time.time() - crafter_step_start
+            logger.debug(f"Crafter env.step() took {crafter_step_time:.3f}s")
+
+            # Update internal state
+            self.obs = obs
+            self.done = done
+            self.info = info
+            self.last_reward = reward
+
+            # Step count is tracked by the crafter environment itself in self.env._step
+
+            # Process achievements - check what was unlocked this step
+            new_achievements = set()
+            if "achievements" in info:
+                for achievement, status in info["achievements"].items():
+                    if status and achievement not in self.achievements_unlocked:
+                        new_achievements.add(achievement)
+                        self.achievements_unlocked.add(achievement)
+
+            # Calculate reward
+            reward_from_stack = 0
+            try:
+                if hasattr(self, "_reward_stack") and self._reward_stack:
+                    reward_from_stack = sum(self._reward_stack)
+                    self._reward_stack.clear()
+            except Exception as e:
+                reward_from_stack = 0
+
+            # Create private state
+            # Current episode reward
+            final_reward = self._total_reward + reward + reward_from_stack
+            self._total_reward = final_reward
+
+            # Determine proper termination reason based on game state
+            player = self.env._player  # type: ignore[attr-defined]
+            current_step = self.env._step  # type: ignore[attr-defined]
+            max_steps = self.env._length  # type: ignore[attr-defined]
+
+            # Check if player died (health <= 0)
+            player_died = player.health <= 0
+
+            # Check if max steps reached
+            max_steps_reached = current_step >= max_steps
+
+            # Set termination flags properly:
+            # - terminated=True only if player actually died
+            # - truncated=True only if episode ended due to step limit
+            if done:
+                if player_died:
+                    terminated = True
+                    truncated = False
+                elif max_steps_reached:
+                    terminated = False
+                    truncated = True
+                else:
+                    # Fallback: if done=True but unclear reason, assume timeout
+                    terminated = False
+                    truncated = True
+            else:
+                terminated = False
+                truncated = False
+
+            final_priv_state = self._build_private_state(final_reward, terminated, truncated)
+
+            # Build public state AFTER step to reflect latest world and achievements
+            pub_state_after = self._build_public_state(obs, info)
+
+            # Store post-step state as baseline for next step
+            self._previous_public_state_for_reward = pub_state_after
+            self._previous_private_state_for_reward = final_priv_state
+
+            total_step_time = time.time() - step_start_time
+            logger.debug(
+                f"CrafterEngine _step_engine took {total_step_time:.3f}s (crafter.step: {crafter_step_time:.3f}s)"
+            )
+            return final_priv_state, pub_state_after
+
+        except Exception as e:
+            # Create error state
+            import traceback
+
+            logger.error(f"Step engine error: {e}")
+            logger.error(traceback.format_exc())
+            error_pub_state = self._get_public_state_from_env()
+            error_pub_state.error_info = f"Step engine error: {e}"
+            error_priv_state = self._get_private_state_from_env(
+                reward=-1.0, terminated=True, truncated=False
+            )
+            return error_priv_state, error_pub_state
+
+    # ------------------------------------------------------------------
+    # Rendering (simple text summary)
+    # ------------------------------------------------------------------
+
+    async def _render(
+        self,
+        private_state: CrafterPrivateState,
+        public_state: CrafterPublicState,
+        get_observation: Optional[GetObservationCallable] = None,
+    ) -> str:  # type: ignore[override]
+        obs_cb = get_observation or CrafterObservationCallable()
+        obs = await obs_cb.get_observation(public_state, private_state)
+        if isinstance(obs, str):
+            return obs
+        if isinstance(obs, dict):
+            header = f"steps: {public_state.num_steps_taken}/{public_state.max_steps_episode} | "
+            header += f"last_r: {private_state.reward_last_step:.2f} | total_r: {private_state.total_reward_episode:.2f}"
+            inv = ", ".join(f"{k}:{v}" for k, v in public_state.inventory.items() if v)
+            ach = ", ".join(k for k, v in public_state.achievements_status.items() if v)
+            return f"{header}\ninv: {inv}\nach: {ach}"
+        return str(obs)
+
+    # ------------------------------------------------------------------
+    # Snapshotting for exact reproducibility
+    # ------------------------------------------------------------------
+
+    async def _serialize_engine(self) -> CrafterEngineSnapshot:
+        world = self.env._world  # type: ignore[attr-defined]
+        objects_state = [None if o is None else serialize_world_object(o) for o in world._objects]
+        # capture total reward and original seed
+        total_reward = self._total_reward
+        snap = CrafterEngineSnapshot(
+            env_raw_state=self.env.save(),
+            total_reward_snapshot=total_reward,
+            crafter_seed=self.env._seed,
+            previous_public_state_snapshot=dataclasses.asdict(
+                self._previous_public_state_for_reward
+            )
+            if self._previous_public_state_for_reward
+            else None,
+            previous_private_state_snapshot=dataclasses.asdict(
+                self._previous_private_state_for_reward
+            )
+            if self._previous_private_state_for_reward
+            else None,
+        )
+        return snap
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snapshot: CrafterEngineSnapshot, task_instance: TaskInstance
+    ) -> "CrafterEngine":
+        engine = cls(task_instance)
+        # Initialize env first to create structures
+        obs = engine.env.reset()
+        # Then load the saved state (this overrides the reset)
+        engine.env.load(snapshot.env_raw_state)
+        engine._total_reward = snapshot.total_reward_snapshot
+        engine.env._seed = snapshot.crafter_seed
+
+        # Initialize engine state attributes that step() expects
+        engine.obs = engine.env.render()
+        engine.done = False
+        engine.info = {}
+        engine.last_reward = 0.0
+
+        # Ensure achievements tracking is initialized
+        engine.achievements_unlocked = set()
+
+        # Re-establish previous states for reward system continuity if first step after load
+        engine._previous_public_state_for_reward = engine._build_public_state(engine.env.render())
+        # Safe comparisons to avoid string vs int errors
+        health_dead = safe_compare(0, engine.env._player.health, ">=")
+        step_exceeded = safe_compare(engine.env._length, engine.env._step, "<=")
+        engine._previous_private_state_for_reward = engine._build_private_state(
+            0.0, health_dead, step_exceeded
+        )
+        return engine
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _build_public_state(
+        self, obs_img: np.ndarray, info: Optional[Dict[str, Any]] | None = None
+    ) -> CrafterPublicState:
+        try:
+            if info is None:
+                player = self.env._player  # type: ignore[attr-defined]
+                # Safe achievement status check
+                achievements_status = {}
+                for k, v in player.achievements.items():
+                    achievements_status[k] = safe_compare(0, v, "<")
+                inventory = player.inventory.copy()
+                semantic = getattr(self.env, "_sem_view", lambda: None)()
+            else:
+                inventory = info.get("inventory", {})
+                # Safe achievement status check from info
+                achievements_status = {}
+                achievements_info = info.get("achievements", {})
+                for k, v in achievements_info.items():
+                    achievements_status[k] = safe_compare(0, v, "<")
+                semantic = info.get("semantic")
+
+            player = self.env._player  # type: ignore[attr-defined]
+            return CrafterPublicState(
+                inventory=inventory,
+                achievements_status=achievements_status,
+                player_position=tuple(player.pos),  # type: ignore[attr-defined]
+                player_direction=player.facing,  # type: ignore[attr-defined]
+                semantic_map=semantic,
+                world_material_map=self.env._world._mat_map.copy(),  # type: ignore[attr-defined]
+                observation_image=obs_img,
+                num_steps_taken=self.env._step,  # type: ignore[attr-defined]
+                max_steps_episode=self.env._length,  # type: ignore[attr-defined]
+                error_info=info.get("error_info") if info else None,
+            )
+        except Exception as e:
+            logging.error(f"Error building public state: {e}")
+            # Return minimal safe state
+            return CrafterPublicState(
+                inventory={},
+                achievements_status={},
+                player_position=(0, 0),
+                player_direction=0,
+                semantic_map=None,
+                world_material_map=np.zeros((1, 1), dtype=np.uint8),
+                observation_image=obs_img
+                if obs_img is not None
+                else np.zeros((64, 64, 3), dtype=np.uint8),
+                num_steps_taken=0,
+                max_steps_episode=10000,
+                error_info=f"State building error: {e}",
+            )
+
+    def _build_private_state(
+        self, reward: float, terminated: bool, truncated: bool
+    ) -> CrafterPrivateState:
+        player = self.env._player  # type: ignore[attr-defined]
+        stats = {
+            "health": player.health,
+            "food": player.inventory.get("food"),
+            "drink": player.inventory.get("drink"),
+            "energy": player.inventory.get("energy"),
+            "_hunger": getattr(player, "_hunger", 0),
+            "_thirst": getattr(player, "_thirst", 0),
+        }
+        return CrafterPrivateState(
+            reward_last_step=reward,
+            total_reward_episode=self._total_reward,
+            achievements_current_values=player.achievements.copy(),
+            terminated=terminated,
+            truncated=truncated,
+            player_internal_stats=stats,
+            world_rng_state_snapshot=self.env._world.random.get_state(),  # type: ignore[attr-defined]
+        )
+
+    def _get_public_state_from_env(self) -> CrafterPublicState:
+        """Helper method to get current public state from synth_ai.environments.environment"""
+        try:
+            obs_img = self.env.render()
+            return self._build_public_state(obs_img)
+        except Exception as e:
+            logging.error(f"Error getting public state from env: {e}")
+            # Return default state
+            return CrafterPublicState(
+                inventory={},
+                achievements_status={},
+                player_position=(0, 0),
+                player_direction=0,
+                semantic_map=None,
+                world_material_map=np.zeros((1, 1), dtype=np.uint8),
+                observation_image=np.zeros((64, 64, 3), dtype=np.uint8),
+                num_steps_taken=0,
+                max_steps_episode=10000,
+                error_info=f"State extraction error: {e}",
+            )
+
+    def _get_private_state_from_env(
+        self, reward: float, terminated: bool, truncated: bool
+    ) -> CrafterPrivateState:
+        """Helper method to get current private state from synth_ai.environments.environment"""
+        try:
+            return self._build_private_state(reward, terminated, truncated)
+        except Exception as e:
+            logging.error(f"Error getting private state from env: {e}")
+            # Return default state
+            return CrafterPrivateState(
+                reward_last_step=reward,
+                total_reward_episode=0.0,
+                achievements_current_values={},
+                terminated=terminated,
+                truncated=truncated,
+                player_internal_stats={},
+                world_rng_state_snapshot=None,
+            )
+
+
+# --- Reward Components ---
+class CrafterAchievementComponent(RewardComponent):
+    async def score(self, state: CrafterPublicState, action: Dict[str, Any]) -> float:
+        prev_achievements = action.get("previous_public_state_achievements", {})
+        current_achievements = state.achievements_status
+        new_achievements = sum(
+            1
+            for ach, status in current_achievements.items()
+            if status and not prev_achievements.get(ach)
+        )
+        return float(new_achievements) * 0.1
+
+
+class CrafterPlayerStatComponent(RewardComponent):
+    async def score(self, state: CrafterPrivateState, action: Dict[str, Any]) -> float:
+        current_health = state.player_internal_stats.get("health", 0)
+        prev_health = action.get("previous_private_state_stats", {}).get("health", current_health)
+        if current_health < prev_health:
+            return -0.05  # Lost health penalty
+        return 0.0
+
+
+class CrafterStepPenaltyComponent(RewardComponent):
+    def __init__(self, penalty: float = -0.001):
+        super().__init__()
+        self.penalty = penalty
+        self.weight = 1.0
+
+    async def score(self, state: Any, action: Any) -> float:
+        return self.penalty

synth_ai/environments/examples/crafter_classic/engine_deterministic_patch.py

@@ -0,0 +1,63 @@
+"""
+Apply once (import this module anywhere before CrafterEngine is used).
+It replaces Env._balance_object so that every per-chunk object list is
+sorted by (x, y, class-name) before any random choice is made – removing
+the hash-based set-iteration nondeterminism that caused the drift.
+"""
+
+import collections
+import crafter
+
+print("[PATCH] Attempting to apply Crafter deterministic patch...")
+
+# -----------------------------------------------------------------------------
+# 1.  Make per–chunk object order stable
+# -----------------------------------------------------------------------------
+if not hasattr(crafter.Env, "_orig_balance_object"):
+    print("[PATCH] Patching crafter.Env._balance_object...")
+    crafter.Env._orig_balance_object = crafter.Env._balance_object
+
+    def _balance_object_det(self, chunk, objs, *args, **kwargs):
+        # cls, material, span_dist, despan_dist, spawn_prob, despawn_prob, ctor, target_fn
+        # were part of the original signature, but *args, **kwargs is more robust.
+        objs = sorted(objs, key=lambda o: (o.pos[0], o.pos[1], o.__class__.__name__))
+        return crafter.Env._orig_balance_object(self, chunk, objs, *args, **kwargs)
+
+    crafter.Env._balance_object = _balance_object_det
+    print("[PATCH] crafter.Env._balance_object patched.")
+else:
+    print("[PATCH] crafter.Env._balance_object already patched or _orig_balance_object exists.")
+
+# -----------------------------------------------------------------------------
+# 2.  Make *chunk* iteration order stable
+# -----------------------------------------------------------------------------
+if not hasattr(crafter.engine.World, "_orig_chunks_prop"):
+    crafter.engine.World._orig_chunks_prop = crafter.engine.World.chunks
+
+    def _chunks_sorted(self):
+        # OrderedDict keeps the sorted key order during iteration
+        return collections.OrderedDict(sorted(self._chunks.items()))
+
+    crafter.engine.World.chunks = property(_chunks_sorted)
+
+# -----------------------------------------------------------------------------
+# 3. NEW: keep per-frame object update order deterministic
+# -----------------------------------------------------------------------------
+if not hasattr(crafter.engine.World, "_orig_objects_prop"):
+    crafter.engine.World._orig_objects_prop = crafter.engine.World.objects  # save original
+
+    @property
+    def _objects_sorted(self):
+        objs = [o for o in self._objects if o]  # Filter out None (removed) objects
+        # stable order: x, y, class-name, creation-index
+        return sorted(
+            objs,
+            key=lambda o: (
+                o.pos[0],
+                o.pos[1],
+                o.__class__.__name__,
+                getattr(o, "_id", 0),
+            ),
+        )
+
+    crafter.engine.World.objects = _objects_sorted

synth_ai/environments/examples/crafter_classic/engine_helpers/action_map.py

@@ -0,0 +1,5 @@
+import crafter.constants as C
+from typing import Dict
+
+# Map each action name to its corresponding index in the crafter package
+CRAFTER_ACTION_MAP: Dict[str, int] = {action_name: idx for idx, action_name in enumerate(C.actions)}