synth-ai 0.2.9.dev2__py3-none-any.whl → 0.2.9.dev4__py3-none-any.whl

examples/warming_up_to_rl/task_app/synth_envs_hosted/storage/__init__.py

@@ -0,0 +1,5 @@
+"""Storage module for Modal Volume operations."""
+
+from .volume import VolumeStorage, storage
+
+__all__ = ["VolumeStorage", "storage"]

examples/warming_up_to_rl/task_app/synth_envs_hosted/storage/volume.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import gzip
+import hashlib
+import json
+import os
+import tarfile
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+class VolumeStorage:
+    """Helpers for Modal Volume storage operations."""
+    
+    def __init__(self, base_path: str = "/data/state") -> None:
+        self.base_path = Path(base_path)
+    
+    def get_snapshot_path(
+        self,
+        rl_run_id: str,
+        kind: str,
+        snapshot_id: str,
+    ) -> Path:
+        """Build the path for a snapshot file."""
+        # Use first 2 chars of snapshot_id for sharding
+        shard1 = snapshot_id[:2] if len(snapshot_id) >= 2 else "00"
+        shard2 = snapshot_id[2:4] if len(snapshot_id) >= 4 else "00"
+        
+        return (
+            self.base_path
+            / "runs"
+            / rl_run_id
+            / kind
+            / shard1
+            / shard2
+            / f"{snapshot_id}.tar.gz"
+        )
+    
+    def get_index_path(self, rl_run_id: str) -> Path:
+        """Get the index file path for a run."""
+        return self.base_path / "runs" / rl_run_id / "index" / "meta.jsonl"
+    
+    def write_snapshot_atomic(
+        self,
+        path: Path,
+        archive_bytes: bytes,
+    ) -> None:
+        """Atomically write a snapshot archive to disk."""
+        # Ensure parent directory exists
+        path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Write to temp file first
+        tmp_path = path.with_suffix(".tmp")
+        with open(tmp_path, "wb") as f:
+            f.write(archive_bytes)
+            f.flush()
+            os.fsync(f.fileno())
+        
+        # Atomic rename
+        os.replace(tmp_path, path)
+    
+    def create_archive(
+        self,
+        state_dict: Dict[str, Any],
+        meta: Dict[str, Any],
+    ) -> bytes:
+        """Create a tar.gz archive with state and metadata."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+            
+            # Write state.json
+            state_path = tmppath / "state.json"
+            with open(state_path, "w") as f:
+                json.dump(state_dict, f, sort_keys=True, indent=2)
+            
+            # Write meta.json
+            meta_path = tmppath / "meta.json"
+            with open(meta_path, "w") as f:
+                json.dump(meta, f, sort_keys=True, indent=2)
+            
+            # Create tar archive
+            tar_path = tmppath / "archive.tar"
+            with tarfile.open(tar_path, "w") as tar:
+                tar.add(state_path, arcname="state.json")
+                tar.add(meta_path, arcname="meta.json")
+            
+            # Compress with gzip
+            with open(tar_path, "rb") as f:
+                tar_bytes = f.read()
+            
+            compressed = gzip.compress(tar_bytes, compresslevel=6)
+            
+            return compressed
+    
+    def extract_archive(self, archive_bytes: bytes) -> tuple[Dict[str, Any], Dict[str, Any]]:
+        """Extract state and metadata from a tar.gz archive."""
+        # Decompress
+        tar_bytes = gzip.decompress(archive_bytes)
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+            
+            # Write tar bytes to temp file
+            tar_path = tmppath / "archive.tar"
+            with open(tar_path, "wb") as f:
+                f.write(tar_bytes)
+            
+            # Extract tar
+            with tarfile.open(tar_path, "r") as tar:
+                tar.extractall(tmppath)
+            
+            # Read state and meta
+            with open(tmppath / "state.json", "r") as f:
+                state = json.load(f)
+            
+            with open(tmppath / "meta.json", "r") as f:
+                meta = json.load(f)
+            
+            return state, meta
+    
+    def compute_snapshot_id(self, archive_bytes: bytes) -> str:
+        """Compute content-addressed snapshot ID."""
+        return hashlib.sha256(archive_bytes).hexdigest()
+    
+    def save_snapshot(
+        self,
+        rl_run_id: str,
+        kind: str,
+        state_dict: Dict[str, Any],
+        config: Optional[Dict[str, Any]] = None,
+        parent_snapshot_id: Optional[str] = None,
+    ) -> tuple[str, str, int]:
+        """Save a snapshot and return (snapshot_id, path, size)."""
+        # Build metadata
+        meta = {
+            "kind": kind,
+            "rl_run_id": rl_run_id,
+            "schema_version": "1.0",
+            "created_at": datetime.utcnow().isoformat(),
+        }
+        
+        if parent_snapshot_id:
+            meta["parent_snapshot_id"] = parent_snapshot_id
+        
+        if config:
+            config_str = json.dumps(config, sort_keys=True)
+            meta["config_hash"] = hashlib.sha256(config_str.encode()).hexdigest()
+        
+        # Create archive
+        archive_bytes = self.create_archive(state_dict, meta)
+        
+        # Compute snapshot ID
+        snapshot_id = self.compute_snapshot_id(archive_bytes)
+        meta["snapshot_id"] = snapshot_id
+        
+        # Recreate archive with snapshot_id in metadata
+        archive_bytes = self.create_archive(state_dict, meta)
+        
+        # Get path and write
+        path = self.get_snapshot_path(rl_run_id, kind, snapshot_id)
+        self.write_snapshot_atomic(path, archive_bytes)
+        
+        # Append to index
+        self.append_to_index(rl_run_id, meta)
+        
+        return snapshot_id, str(path), len(archive_bytes)
+    
+    def load_snapshot(
+        self,
+        rl_run_id: str,
+        kind: str,
+        snapshot_id: str,
+    ) -> tuple[Dict[str, Any], Dict[str, Any]]:
+        """Load a snapshot and return (state_dict, meta)."""
+        path = self.get_snapshot_path(rl_run_id, kind, snapshot_id)
+        
+        if not path.exists():
+            raise FileNotFoundError(f"Snapshot not found: {path}")
+        
+        with open(path, "rb") as f:
+            archive_bytes = f.read()
+        
+        state, meta = self.extract_archive(archive_bytes)
+        return state, meta
+    
+    def append_to_index(
+        self,
+        rl_run_id: str,
+        meta: Dict[str, Any],
+    ) -> None:
+        """Append metadata to the run's index file."""
+        index_path = self.get_index_path(rl_run_id)
+        index_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(index_path, "a") as f:
+            f.write(json.dumps(meta) + "\n")
+    
+    def read_index(self, rl_run_id: str) -> list[Dict[str, Any]]:
+        """Read all entries from a run's index file."""
+        index_path = self.get_index_path(rl_run_id)
+        
+        if not index_path.exists():
+            return []
+        
+        entries = []
+        with open(index_path, "r") as f:
+            for line in f:
+                if line.strip():
+                    entries.append(json.loads(line))
+        
+        return entries
+
+
+# Global storage instance
+storage = VolumeStorage()

examples/warming_up_to_rl/task_app/synth_envs_hosted/test_agents.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""
+Smoke test for Wordle and Sokoban ReAct agents using the hosted service.
+
+Prereqs:
+- Run the service: python backend/app/services/rl/online/synth_envs_hosted/main.py
+- Run an OpenAI-compatible inference server (e.g., Flash/vLLM) at VLLM_BASE_URL
+  that serves model "gpt-5-nano" or adjust MODEL below.
+
+This script will:
+- Create a Wordle/Sokoban env
+- Create corresponding *-react policy with tools
+- Ask the policy for tool_calls via /policy/step (which calls the model)
+- Apply tool_calls to the env via /env/step
+"""
+
+import asyncio
+import os
+import httpx
+
+BASE_URL = os.environ.get("SYNTH_ENVS_HOSTED_URL", "http://localhost:8000")
+INFER_URL = os.environ.get("VLLM_BASE_URL", "http://localhost:8001")
+MODEL = os.environ.get("MODEL", "gpt-5-nano")
+
+
+async def run_wordle(rounds: int = 3) -> None:
+    async with httpx.AsyncClient() as client:
+        # Create env
+        resp = await client.post(
+            f"{BASE_URL}/env/create",
+            json={
+                "env_name": "Wordle",
+                "config": {"word_length": 5, "max_guesses": 6},
+                "seed": 0,
+                "rl_run_id": "agents-smoke",
+            },
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        env_id = data["env_id"]
+        obs = data["observation"]
+        print("Wordle env created:", env_id)
+
+        # Create policy
+        resp = await client.post(
+            f"{BASE_URL}/policy/create",
+            json={
+                "policy_name": "wordle-react",
+                "config": {
+                    "inference_url": INFER_URL,
+                    "model": MODEL,
+                    "use_tools": True,
+                    "word_length": 5,
+                    "max_guesses": 6,
+                },
+                "rl_run_id": "agents-smoke",
+                "bound_env_id": env_id,
+            },
+        )
+        resp.raise_for_status()
+        policy_id = resp.json()["policy_id"]
+        print("Wordle policy:", policy_id)
+
+        # Loop a few rounds
+        for i in range(rounds):
+            print(f"[Wordle] Round {i + 1}")
+            step_req = {"policy_id": policy_id, "observation": obs, "dry_run": False}
+            resp = await client.post(f"{BASE_URL}/policy/step", json=step_req)
+            resp.raise_for_status()
+            step_out = resp.json()
+            tool_calls = step_out.get("tool_calls", [])
+            print(" tool_calls:", tool_calls)
+            if not tool_calls:
+                break
+            resp = await client.post(
+                f"{BASE_URL}/env/step",
+                json={"env_id": env_id, "tool_calls": tool_calls},
+            )
+            resp.raise_for_status()
+            env_step = resp.json()
+            obs = env_step["observation"]
+            print(" done:", env_step.get("done"), "reward:", env_step.get("reward"))
+            if env_step.get("done"):
+                break
+
+
+async def run_sokoban(rounds: int = 3) -> None:
+    async with httpx.AsyncClient() as client:
+        # Create env (no initial_state provided; relies on env default)
+        resp = await client.post(
+            f"{BASE_URL}/env/create",
+            json={
+                "env_name": "Sokoban",
+                "config": {"difficulty": "easy"},
+                "seed": 0,
+                "rl_run_id": "agents-smoke",
+            },
+        )
+        if resp.status_code != 200:
+            print("Sokoban create failed:", resp.status_code, resp.text)
+            return
+        data = resp.json()
+        env_id = data["env_id"]
+        obs = data["observation"]
+        print("Sokoban env created:", env_id)
+
+        resp = await client.post(
+            f"{BASE_URL}/policy/create",
+            json={
+                "policy_name": "sokoban-react",
+                "config": {
+                    "inference_url": INFER_URL,
+                    "model": MODEL,
+                    "use_tools": True,
+                },
+                "rl_run_id": "agents-smoke",
+                "bound_env_id": env_id,
+            },
+        )
+        if resp.status_code != 200:
+            print("Sokoban policy create failed:", resp.status_code, resp.text)
+            return
+        policy_id = resp.json()["policy_id"]
+        print("Sokoban policy:", policy_id)
+
+        for i in range(rounds):
+            print(f"[Sokoban] Round {i + 1}")
+            step_req = {"policy_id": policy_id, "observation": obs, "dry_run": False}
+            resp = await client.post(f"{BASE_URL}/policy/step", json=step_req)
+            if resp.status_code != 200:
+                print(" policy step failed:", resp.status_code, resp.text)
+                break
+            step_out = resp.json()
+            tool_calls = step_out.get("tool_calls", [])
+            print(" tool_calls:", tool_calls)
+            if not tool_calls:
+                break
+            resp = await client.post(
+                f"{BASE_URL}/env/step",
+                json={"env_id": env_id, "tool_calls": tool_calls},
+            )
+            if resp.status_code != 200:
+                print(" env step failed:", resp.status_code, resp.text)
+                break
+            env_step = resp.json()
+            obs = env_step["observation"]
+            print(" done:", env_step.get("done"), "reward:", env_step.get("reward"))
+            if env_step.get("done"):
+                break
+
+
+async def main():
+    print("Testing Wordle agent with model:", MODEL)
+    await run_wordle(rounds=3)
+    print("\nTesting Sokoban agent with model:", MODEL)
+    await run_sokoban(rounds=3)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/warming_up_to_rl/task_app/synth_envs_hosted/test_service.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""
+Simple test script for the GRPO Synth Envs Hosted Service.
+
+Run this after starting the service with:
+    python main.py
+"""
+
+import asyncio
+import json
+import httpx
+
+
+async def test_service():
+    """Test basic service functionality."""
+    base_url = "http://localhost:8000"
+
+    async with httpx.AsyncClient() as client:
+        # Test 1: Service info
+        print("1. Testing /info endpoint...")
+        response = await client.get(f"{base_url}/info")
+        assert response.status_code == 200
+        info = response.json()
+        print(f"   Service info: {json.dumps(info, indent=2)}")
+
+        # Test 2: Health check
+        print("\n2. Testing /health endpoint...")
+        response = await client.get(f"{base_url}/health")
+        assert response.status_code == 200
+        print(f"   Health: {response.json()}")
+
+        # Test 3: Create environment
+        print("\n3. Creating environment...")
+        response = await client.post(
+            f"{base_url}/env/create",
+            json={
+                "env_name": "crafter",
+                "config": {},
+                "seed": 42,
+                "rl_run_id": "test-run-001",
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+            return
+        env_data = response.json()
+        env_id = env_data["env_id"]
+        print(f"   Created env: {env_id}")
+        print(f"   Initial observation keys: {list(env_data['observation'].keys())}")
+
+        # Test 4: Create policy
+        print("\n4. Creating policy...")
+        response = await client.post(
+            f"{base_url}/policy/create",
+            json={
+                "policy_name": "crafter-react",
+                "config": {
+                    "inference_url": "http://localhost:8001",
+                    "model": "test-model",
+                },
+                "rl_run_id": "test-run-001",
+                "bound_env_id": env_id,
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+            return
+        policy_data = response.json()
+        policy_id = policy_data["policy_id"]
+        print(f"   Created policy: {policy_id}")
+
+        # Test 5: Environment step with dummy tool calls
+        print("\n5. Testing environment step...")
+        response = await client.post(
+            f"{base_url}/env/step",
+            json={
+                "env_id": env_id,
+                "tool_calls": [{"tool": "interact", "args": {"action": "move_left"}}],
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            step_data = response.json()
+            print(
+                f"   Step result - done: {step_data['done']}, reward: {step_data.get('reward')}"
+            )
+
+        # Test 6: Environment snapshot
+        print("\n6. Creating environment snapshot...")
+        response = await client.post(
+            f"{base_url}/env/snapshot", json={"env_id": env_id}
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            snapshot_data = response.json()
+            print(f"   Snapshot ID: {snapshot_data['snapshot_id']}")
+            print(f"   Size: {snapshot_data['size']} bytes")
+
+        # Test 7: Policy snapshot
+        print("\n7. Creating policy snapshot...")
+        response = await client.post(
+            f"{base_url}/policy/snapshot", json={"policy_id": policy_id}
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            snapshot_data = response.json()
+            print(f"   Snapshot ID: {snapshot_data['snapshot_id']}")
+            print(f"   Size: {snapshot_data['size']} bytes")
+
+        # Test 8: Run status
+        print("\n8. Testing run status...")
+        response = await client.get(f"{base_url}/run/status/test-run-001")
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            status_data = response.json()
+            print(f"   Run status: {status_data['status']}")
+
+        # Test 9: Terminate environment
+        print("\n9. Terminating environment...")
+        response = await client.post(
+            f"{base_url}/env/terminate", json={"env_id": env_id}
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            print(f"   Environment terminated: {response.json()['ok']}")
+
+        # Test 10: Terminate policy
+        print("\n10. Terminating policy...")
+        response = await client.post(
+            f"{base_url}/policy/terminate", json={"policy_id": policy_id}
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            print(f"   Policy terminated: {response.json()['ok']}")
+
+        print("\n✅ All basic tests completed!")
+
+
+if __name__ == "__main__":
+    asyncio.run(test_service())

examples/warming_up_to_rl/task_app/synth_envs_hosted/test_stepwise_rewards.py

@@ -0,0 +1,58 @@
+import math
+
+from backend.app.routes.clustered_training.dev.synth_envs_hosted.rollout import (
+    compute_stepwise_reward,
+)
+
+
+def test_compute_stepwise_reward_detects_unlock() -> None:
+    prev = {"build_workbench": False, "pet_cow": False}
+    new = {"build_workbench": True, "pet_cow": False}
+    actions = [{"tool": "interact", "args": {"action": "build_workbench"}}]
+
+    info, decision, stats = compute_stepwise_reward(
+        prev,
+        new,
+        decision_index=2,
+        actions_summary=actions,
+        indicator_lambda=0.75,
+    )
+
+    assert info["decision_index"] == 2
+    assert info["new_achievements"] == ["build_workbench"]
+    assert info["indicator"] == 1
+    assert math.isclose(info["reward"], 0.75)
+
+    assert decision["actions"] == actions
+    assert decision["indicator"] == 1
+    assert math.isclose(decision["r_i"], 0.75)
+
+    assert math.isclose(stats["indicator"], 1.0)
+    assert math.isclose(stats["reward"], 0.75)
+    assert math.isclose(stats["new_achievements_count"], 1.0)
+
+
+def test_compute_stepwise_reward_handles_no_unlock() -> None:
+    prev = {"collect_wheat": True, "milk_cow": False}
+    new = {"collect_wheat": True, "milk_cow": False}
+    actions = [{"tool": "interact", "args": {"action": "sleep"}}]
+
+    info, decision, stats = compute_stepwise_reward(
+        prev,
+        new,
+        decision_index=5,
+        actions_summary=actions,
+        indicator_lambda=1.2,
+    )
+
+    assert info["new_achievements"] == []
+    assert info["indicator"] == 0
+    assert math.isclose(info["reward"], 0.0)
+
+    assert decision["actions"] == actions
+    assert decision["indicator"] == 0
+    assert math.isclose(decision["r_i"], 0.0)
+
+    assert math.isclose(stats["indicator"], 0.0)
+    assert math.isclose(stats["reward"], 0.0)
+    assert math.isclose(stats["new_achievements_count"], 0.0)

examples/warming_up_to_rl/task_app/synth_envs_hosted/utils.py

@@ -0,0 +1,61 @@
+"""Utility functions for the task service."""
+
+import numpy as np
+from typing import Any, Dict, List, Union
+
+
+def convert_numpy_to_python(obj: Any) -> Any:
+    """
+    Recursively convert numpy types to Python native types for JSON serialization.
+
+    Args:
+        obj: Object that may contain numpy types
+
+    Returns:
+        Object with numpy types converted to Python native types
+    """
+    if isinstance(obj, np.integer):
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        return float(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    elif isinstance(obj, dict):
+        return {key: convert_numpy_to_python(value) for key, value in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [convert_numpy_to_python(item) for item in obj]
+    else:
+        return obj
+
+
+def sanitize_observation(observation: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Sanitize observation data for JSON serialization.
+
+    Converts numpy types and removes non-serializable objects.
+
+    Args:
+        observation: Raw observation from environment
+
+    Returns:
+        Sanitized observation safe for JSON serialization
+    """
+    if not isinstance(observation, dict):
+        return observation
+
+    sanitized = {}
+    for key, value in observation.items():
+        # Skip non-serializable keys or convert them
+        if key in ["semantic_map", "world_material_map", "observation_image"]:
+            # These are likely numpy arrays - convert to lists or skip
+            if isinstance(value, np.ndarray):
+                # For large arrays, we might want to skip or compress
+                # For now, skip them as they're likely debug info
+                continue
+        elif key == "player_position" and isinstance(value, tuple):
+            # Convert tuple with potential numpy types
+            sanitized[key] = [convert_numpy_to_python(v) for v in value]
+        else:
+            sanitized[key] = convert_numpy_to_python(value)
+
+    return sanitized