synth-ai 0.2.0__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/examples/nethack/taskset.py

@@ -0,0 +1,323 @@
+"""TaskSet generation for NetHack environment."""
+
+from __future__ import annotations
+
+import random
+from uuid import uuid4
+from dataclasses import dataclass
+from typing import Dict, Any, List, Optional, Set
+
+from synth_ai.environments.tasks.core import (
+    TaskInstance,
+    TaskInstanceMetadata,
+    TaskInstanceSet,
+    Impetus,
+    Intent,
+    SplitInfo,
+)
+
+
+@dataclass
+class NetHackTaskInstanceMetadata(TaskInstanceMetadata):
+    """Task-specific metadata for NetHack."""
+
+    character_role: str  # "wizard", "knight", etc.
+    starting_level: int  # Dungeon level to start on
+    target_depth: int  # Goal depth to reach
+    time_limit: int  # Maximum turns
+    difficulty: str  # "easy", "medium", "hard"
+    special_objectives: List[str]  # Additional goals beyond survival
+    seed: Optional[int] = None  # Random seed for reproducibility
+
+
+@dataclass
+class NetHackTaskInstance(TaskInstance):
+    """NetHack task instance."""
+
+    async def serialize(self) -> dict:
+        """Convert to serializable format."""
+        return {
+            "id": str(self.id),
+            "impetus": {"instructions": self.impetus.instructions},
+            "intent": {
+                "rubric": self.intent.rubric,
+                "gold_trajectories": None,
+                "gold_state_diff": self.intent.gold_state_diff,
+            },
+            "metadata": {
+                "character_role": self.metadata.character_role,
+                "starting_level": self.metadata.starting_level,
+                "target_depth": self.metadata.target_depth,
+                "time_limit": self.metadata.time_limit,
+                "difficulty": self.metadata.difficulty,
+                "special_objectives": self.metadata.special_objectives,
+                "seed": self.metadata.seed,
+            },
+            "is_reproducible": self.is_reproducible,
+            "initial_engine_snapshot": None,
+        }
+
+    @classmethod
+    async def deserialize(cls, data: dict) -> "NetHackTaskInstance":
+        """Restore from serialized data."""
+        return cls(
+            id=uuid4() if not data.get("id") else data["id"],
+            impetus=Impetus(instructions=data["impetus"]["instructions"]),
+            intent=Intent(
+                rubric=data["intent"]["rubric"],
+                gold_trajectories=None,
+                gold_state_diff=data["intent"]["gold_state_diff"],
+            ),
+            metadata=NetHackTaskInstanceMetadata(
+                character_role=data["metadata"]["character_role"],
+                starting_level=data["metadata"]["starting_level"],
+                target_depth=data["metadata"]["target_depth"],
+                time_limit=data["metadata"]["time_limit"],
+                difficulty=data["metadata"]["difficulty"],
+                special_objectives=data["metadata"]["special_objectives"],
+                seed=data["metadata"].get("seed"),
+            ),
+            is_reproducible=data.get("is_reproducible", True),
+            initial_engine_snapshot=None,
+        )
+
+
+# Character role definitions
+CHARACTER_ROLES = {
+    "tourist": {
+        "description": "A tourist with a camera and Hawaiian shirt",
+        "difficulty_modifier": 0.8,  # Easier
+        "starting_items": ["camera", "credit card", "hawaiian shirt"],
+        "strengths": ["gold finding", "luck"],
+        "weaknesses": ["combat", "magic"],
+    },
+    "knight": {
+        "description": "A noble knight in shining armor",
+        "difficulty_modifier": 1.0,
+        "starting_items": ["long sword", "armor", "shield"],
+        "strengths": ["combat", "riding"],
+        "weaknesses": ["magic"],
+    },
+    "wizard": {
+        "description": "A powerful wizard with magical abilities",
+        "difficulty_modifier": 1.2,
+        "starting_items": ["quarterstaff", "spellbook", "cloak"],
+        "strengths": ["magic", "identify"],
+        "weaknesses": ["physical combat", "low hp"],
+    },
+    "barbarian": {
+        "description": "A fierce barbarian warrior",
+        "difficulty_modifier": 0.9,
+        "starting_items": ["battle axe", "leather armor"],
+        "strengths": ["combat", "hp", "strength"],
+        "weaknesses": ["magic", "intelligence"],
+    },
+    "ranger": {
+        "description": "A skilled ranger and tracker",
+        "difficulty_modifier": 1.0,
+        "starting_items": ["bow", "arrows", "cloak"],
+        "strengths": ["ranged combat", "stealth"],
+        "weaknesses": ["melee combat"],
+    },
+    "priest": {
+        "description": "A holy priest with divine powers",
+        "difficulty_modifier": 1.1,
+        "starting_items": ["mace", "robe", "holy water"],
+        "strengths": ["healing", "undead turning"],
+        "weaknesses": ["edged weapons"],
+    },
+    "monk": {
+        "description": "A disciplined monk with martial arts skills",
+        "difficulty_modifier": 1.3,
+        "starting_items": ["robe"],
+        "strengths": ["martial arts", "speed"],
+        "weaknesses": ["armor", "weapons"],
+    },
+    "rogue": {
+        "description": "A stealthy rogue and thief",
+        "difficulty_modifier": 1.1,
+        "starting_items": ["dagger", "leather armor", "lock pick"],
+        "strengths": ["stealth", "backstab", "traps"],
+        "weaknesses": ["direct combat"],
+    },
+}
+
+# Special objectives for variety
+SPECIAL_OBJECTIVES = {
+    "exploration": [
+        "Explore at least 3 different dungeon levels",
+        "Find and enter a shop",
+        "Discover a special room (vault, zoo, etc.)",
+        "Find the entrance to the Gnomish Mines",
+    ],
+    "combat": [
+        "Defeat 10 monsters",
+        "Defeat a monster using magic",
+        "Defeat a monster using ranged weapons",
+        "Survive an encounter with a tough monster",
+    ],
+    "collection": [
+        "Collect 100 gold pieces",
+        "Find and identify a magical item",
+        "Collect food rations for survival",
+        "Find a valuable gem",
+    ],
+    "survival": [
+        "Survive for 500 turns",
+        "Maintain full health for 100 turns",
+        "Never let hunger status reach 'Weak'",
+        "Avoid all traps",
+    ],
+    "progression": [
+        "Reach experience level 3",
+        "Improve at least one skill",
+        "Successfully pray to your deity",
+        "Complete a quest or mission",
+    ],
+}
+
+
+async def create_nethack_taskset() -> TaskInstanceSet:
+    """Generate diverse NetHack scenarios."""
+    instances = []
+
+    # Configuration for different difficulty levels
+    DIFFICULTY_CONFIGS = {
+        "tutorial": {
+            "roles": ["tourist"],
+            "target_depth_range": (1, 3),
+            "time_limit_range": (500, 1000),
+            "objective_count": 1,
+            "count": 20,
+        },
+        "beginner": {
+            "roles": ["knight", "barbarian"],
+            "target_depth_range": (3, 5),
+            "time_limit_range": (1000, 2000),
+            "objective_count": 2,
+            "count": 30,
+        },
+        "intermediate": {
+            "roles": ["wizard", "ranger", "priest"],
+            "target_depth_range": (5, 10),
+            "time_limit_range": (2000, 5000),
+            "objective_count": 2,
+            "count": 25,
+        },
+        "advanced": {
+            "roles": ["monk", "rogue"],
+            "target_depth_range": (10, 15),
+            "time_limit_range": (5000, 10000),
+            "objective_count": 3,
+            "count": 15,
+        },
+        "expert": {
+            "roles": list(CHARACTER_ROLES.keys()),
+            "target_depth_range": (15, 20),
+            "time_limit_range": (10000, 20000),
+            "objective_count": 4,
+            "count": 10,
+        },
+    }
+
+    # Generate instances for each difficulty
+    for difficulty, config in DIFFICULTY_CONFIGS.items():
+        for i in range(config["count"]):
+            # Random role selection
+            role = random.choice(config["roles"])
+            role_info = CHARACTER_ROLES[role]
+
+            # Random parameters within difficulty range
+            min_depth, max_depth = config["target_depth_range"]
+            target_depth = random.randint(min_depth, max_depth)
+            min_time, max_time = config["time_limit_range"]
+            time_limit = random.randint(min_time, max_time)
+
+            # Select random objectives
+            objectives = []
+            objective_categories = list(SPECIAL_OBJECTIVES.keys())
+            for _ in range(config["objective_count"]):
+                category = random.choice(objective_categories)
+                objective = random.choice(SPECIAL_OBJECTIVES[category])
+                objectives.append(objective)
+
+            # Create instruction text
+            instructions = f"""You are a {role_info["description"]}.
+
+Your primary goal is to descend to dungeon level {target_depth} within {time_limit} turns.
+
+Additional objectives:
+{chr(10).join(f"- {obj}" for obj in objectives)}
+
+Character strengths: {", ".join(role_info["strengths"])}
+Character weaknesses: {", ".join(role_info["weaknesses"])}
+
+Tips:
+- Use 'inventory' to check your items
+- Use 'search' to find secret doors
+- Eat food before you become weak from hunger
+- Save valuable items for when you need them
+- Be cautious around unfamiliar monsters
+
+Remember: In NetHack, careful planning often beats hasty action!"""
+
+            # Create success criteria
+            rubric = {
+                "goal": f"Reach dungeon level {target_depth}",
+                "success_criteria": {
+                    "primary": f"Reach dungeon level {target_depth} within {time_limit} turns",
+                    "secondary": objectives,
+                },
+                "evaluation_metrics": {
+                    "depth_reached": target_depth,
+                    "time_limit": time_limit,
+                    "objectives_completed": len(objectives),
+                },
+            }
+
+            # Create metadata
+            metadata = NetHackTaskInstanceMetadata(
+                character_role=role,
+                starting_level=1,
+                target_depth=target_depth,
+                time_limit=time_limit,
+                difficulty=difficulty,
+                special_objectives=objectives,
+                seed=random.randint(0, 2**31 - 1),
+            )
+
+            # Create task instance
+            instance = NetHackTaskInstance(
+                id=uuid4(),
+                impetus=Impetus(instructions=instructions),
+                intent=Intent(rubric=rubric, gold_trajectories=None, gold_state_diff={}),
+                metadata=metadata,
+                is_reproducible=True,
+                initial_engine_snapshot=None,
+            )
+
+            instances.append(instance)
+
+    # Define splits (80% train, 10% val, 10% test)
+    random.shuffle(instances)
+    n_instances = len(instances)
+    n_val = n_instances // 10
+    n_test = n_instances // 10
+
+    val_ids = {inst.id for inst in instances[:n_val]}
+    test_ids = {inst.id for inst in instances[n_val : n_val + n_test]}
+
+    split_info = SplitInfo(
+        val_instance_ids=val_ids, test_instance_ids=test_ids, _is_split_defined=True
+    )
+
+    return TaskInstanceSet(
+        name="NetHack TaskSet",
+        description="A comprehensive set of NetHack dungeon exploration tasks with varying difficulty levels, character roles, and objectives",
+        instances=instances,
+        split_info=split_info,
+    )
+
+
+# Module-level export
+taskset = create_nethack_taskset

synth_ai/environments/examples/nethack/units/test_nethack_engine.py

@@ -0,0 +1,277 @@
+"""Unit tests for NetHack engine."""
+
+import pytest
+import asyncio
+from uuid import uuid4
+
+from synth_ai.environments.tasks.core import TaskInstance, Impetus, Intent, TaskInstanceMetadata
+
+from synth_ai.environments.examples.nethack.engine import (
+    NetHackEngine,
+    NetHackPublicState,
+    NetHackPrivateState,
+    NetHackEngineSnapshot,
+    NetHackSurvivalComponent,
+    NetHackProgressComponent,
+)
+from synth_ai.environments.examples.nethack.taskset import (
+    NetHackTaskInstanceMetadata,
+    NetHackTaskInstance,
+)
+
+# Since engine requires NLE, all tests require it
+pytest.importorskip("nle", reason="NLE is required for NetHack engine")
+
+
+class TestNetHackEngine:
+    """Test cases for NetHack engine."""
+
+    @pytest.fixture
+    def mock_task_instance(self):
+        """Create a mock task instance for testing."""
+        metadata = NetHackTaskInstanceMetadata(
+            character_role="tourist",
+            starting_level=1,
+            target_depth=5,
+            time_limit=1000,
+            difficulty="beginner",
+            special_objectives=["Collect 100 gold pieces"],
+            seed=42,
+        )
+
+        return NetHackTaskInstance(
+            id=uuid4(),
+            impetus=Impetus(instructions="Test NetHack game"),
+            intent=Intent(
+                rubric={"goal": "Reach depth 5"},
+                gold_trajectories=None,
+                gold_state_diff={},
+            ),
+            metadata=metadata,
+            is_reproducible=True,
+            initial_engine_snapshot=None,
+        )
+
+    @pytest.mark.asyncio
+    async def test_engine_initialization(self, mock_task_instance):
+        """Test engine initialization."""
+        engine = NetHackEngine(mock_task_instance)
+
+        assert engine.task_instance == mock_task_instance
+        assert engine.character_role == "tourist"
+        assert engine.max_turns == 1000
+        assert engine.public_state is None
+        assert engine.private_state is None
+
+        # Cleanup
+        if hasattr(engine, "nle"):
+            engine.nle.close()
+
+    @pytest.mark.asyncio
+    async def test_engine_reset(self, mock_task_instance):
+        """Test engine reset functionality."""
+        engine = NetHackEngine(mock_task_instance)
+
+        priv, pub = await engine._reset_engine(seed=42)
+
+        # Check private state
+        assert isinstance(priv, NetHackPrivateState)
+        assert priv.reward_last == 0.0
+        assert priv.total_reward == 0.0
+        assert priv.terminated is False
+        assert priv.score >= 0
+        assert priv.depth_reached >= 1
+
+        # Check public state
+        assert isinstance(pub, NetHackPublicState)
+        assert pub.dungeon_level >= 1
+        assert pub.terminated is False
+        assert pub.turn_count == 0
+        assert pub.max_turns == 1000
+        assert len(pub.message) > 0  # Should have some message
+        assert len(pub.ascii_map) > 0  # Should have a map
+
+    @pytest.mark.asyncio
+    async def test_basic_movement(self, mock_task_instance):
+        """Test basic movement actions."""
+        engine = NetHackEngine(mock_task_instance)
+        priv0, pub0 = await engine._reset_engine()
+        initial_pos = pub0.position
+
+        # Test wait action (always valid)
+        priv, pub = await engine._step_engine("wait")
+        assert pub.last_action == "wait"
+        assert pub.turn_count == 1
+        assert priv.reward_last >= 0  # Should get at least survival reward
+
+    @pytest.mark.asyncio
+    async def test_invalid_action(self, mock_task_instance):
+        """Test invalid action handling."""
+        engine = NetHackEngine(mock_task_instance)
+        await engine._reset_engine()
+
+        # Test invalid action
+        with pytest.raises(ValueError, match="Invalid action|Valid actions"):
+            await engine._step_engine("invalid_action_xyz")
+
+    @pytest.mark.asyncio
+    async def test_turn_limit(self, mock_task_instance):
+        """Test turn limit enforcement."""
+        # Create instance with very short time limit
+        metadata = NetHackTaskInstanceMetadata(
+            character_role="tourist",
+            starting_level=1,
+            target_depth=5,
+            time_limit=3,  # Very short
+            difficulty="test",
+            special_objectives=[],
+            seed=42,
+        )
+
+        task = NetHackTaskInstance(
+            id=uuid4(),
+            impetus=Impetus(instructions="Test"),
+            intent=Intent(rubric={"goal": "Test"}, gold_trajectories=None, gold_state_diff={}),
+            metadata=metadata,
+            is_reproducible=True,
+            initial_engine_snapshot=None,
+        )
+
+        engine = NetHackEngine(task)
+        await engine._reset_engine()
+
+        # Take actions until time limit
+        await engine._step_engine("wait")
+        await engine._step_engine("wait")
+        priv, pub = await engine._step_engine("wait")
+
+        # Should be terminated due to time limit
+        assert pub.terminated is True
+        assert priv.terminated is True
+        assert priv.truncated is True
+        assert "Time limit reached" in pub.message
+
+    @pytest.mark.asyncio
+    async def test_reward_calculation(self, mock_task_instance):
+        """Test reward calculation."""
+        engine = NetHackEngine(mock_task_instance)
+        await engine._reset_engine()
+
+        # Test survival reward
+        priv1, pub1 = await engine._step_engine("wait")
+        # Rewards can be positive (survival) or negative (NLE penalty)
+        assert priv1.reward_last != 0  # Should get some reward
+
+        # Test multiple steps accumulate reward
+        total_before = priv1.total_reward
+        priv2, pub2 = await engine._step_engine("wait")
+        # Total reward should change
+        assert priv2.total_reward != total_before
+
+    @pytest.mark.asyncio
+    async def test_state_diff(self, mock_task_instance):
+        """Test state diff functionality."""
+        engine = NetHackEngine(mock_task_instance)
+        priv1, pub1 = await engine._reset_engine()
+
+        # Since the engine returns references to its internal state,
+        # we need to capture the values we want to compare
+        initial_turn_count = pub1.turn_count
+        initial_last_action = pub1.last_action
+        initial_reward_last = priv1.reward_last
+        initial_total_reward = priv1.total_reward
+
+        # Take an action
+        priv2, pub2 = await engine._step_engine("wait")
+
+        # Verify states actually changed
+        assert pub2.turn_count == initial_turn_count + 1
+        assert pub2.last_action == "wait"
+        assert priv2.total_reward > initial_total_reward  # Should have some reward
+
+        # Since pub1 and pub2 are the same object (engine returns references),
+        # the diff will be empty. This is expected behavior for a stateful engine.
+        # The test was incorrectly assuming the engine returns copies.
+
+    @pytest.mark.asyncio
+    async def test_serialization_roundtrip(self, mock_task_instance):
+        """Test state serialization and deserialization."""
+        engine = NetHackEngine(mock_task_instance)
+        await engine._reset_engine()
+
+        # Take some actions
+        await engine._step_engine("wait")
+        await engine._step_engine("search")
+
+        # Serialize
+        snapshot = await engine._serialize_engine()
+        assert isinstance(snapshot, NetHackEngineSnapshot)
+
+        # Deserialize
+        restored_engine = await NetHackEngine._deserialize_engine(snapshot)
+
+        # Check restored state matches
+        orig_priv, orig_pub = engine.get_current_states_for_observation()
+        rest_priv, rest_pub = restored_engine.get_current_states_for_observation()
+
+        assert rest_pub.turn_count == orig_pub.turn_count
+        assert rest_pub.position == orig_pub.position
+        assert rest_priv.total_reward == orig_priv.total_reward
+
+    @pytest.mark.asyncio
+    async def test_get_current_states(self, mock_task_instance):
+        """Test getting current states without advancing."""
+        engine = NetHackEngine(mock_task_instance)
+
+        # Should raise before initialization
+        with pytest.raises(RuntimeError, match="Engine not initialized"):
+            engine.get_current_states_for_observation()
+
+        # Initialize
+        await engine._reset_engine()
+
+        # Get states multiple times
+        priv1, pub1 = engine.get_current_states_for_observation()
+        priv2, pub2 = engine.get_current_states_for_observation()
+
+        # Should be same states
+        assert pub1.turn_count == pub2.turn_count
+        assert priv1.total_reward == priv2.total_reward
+
+
+class TestRewardComponents:
+    """Test reward components."""
+
+    @pytest.mark.asyncio
+    async def test_survival_component(self):
+        """Test survival reward component."""
+        component = NetHackSurvivalComponent()
+
+        # Test alive state
+        state = NetHackPublicState(terminated=False)
+        reward = await component.score(state, "wait")
+        assert reward == 0.01
+
+        # Test dead state
+        state.terminated = True
+        reward = await component.score(state, "wait")
+        assert reward == -1.0
+
+    @pytest.mark.asyncio
+    async def test_progress_component(self):
+        """Test progress reward component."""
+        component = NetHackProgressComponent()
+
+        # Test no progress
+        state = NetHackPublicState(dungeon_level=1)
+        reward = await component.score(state, "wait")
+        assert reward == 0.0
+
+        # Test depth increase
+        state.dungeon_level = 2
+        reward = await component.score(state, "go_down")
+        assert reward == 1.0
+
+        # Test no further reward for same level
+        reward = await component.score(state, "wait")
+        assert reward == 0.0