synth-ai 0.2.9.dev0__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

examples/task_apps/crafter/task_app/synth_envs_hosted/envs/crafter/policy.py

@@ -0,0 +1,583 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from .react_agent import CrafterReActAgent
+from .tools import TOOLS_SCHEMA
+
+
+# Define Policy base class here to avoid circular import
+class Policy(ABC):
+    """Base class for environment-specific policies."""
+
+    @abstractmethod
+    def prepare_inference_request(
+        self, observation: dict[str, Any], history: list[dict[str, Any]] = None
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]:
+        """Prepare an inference request."""
+        pass
+
+    @abstractmethod
+    def parse_model_response(
+        self, response: str, observation: dict[str, Any]
+    ) -> list[dict[str, Any]]:
+        """Parse model response into tool calls."""
+        pass
+
+
+# (imports moved to top of file to satisfy linter)
+
+
+class CrafterPolicy(Policy):
+    """Thin policy scaffold for Crafter using the ReAct agent prompts.
+
+    This class does not run inference itself. It prepares an inference request
+    (messages and optional tools schema) that the Task App can send to the
+    inference service, and provides helpers to parse the model response into
+    environment tool calls.
+    """
+
+    name: str = "crafter-react"
+
+    def __init__(self, inference_url: str, model: str | None = None) -> None:
+        self.inference_url = inference_url
+        self.model = model
+        self.use_tools = True
+        self.use_vision = False  # Enable vision for VLMs
+        self.image_only_mode = False  # If True, only send images without text observations
+        # Sampling parameters (populated via initialize(config))
+        self.temperature: float | None = None
+        self.top_p: float | None = None
+        self.max_tokens: int | None = None
+        # Thinking controls (populated via initialize(config))
+        self.thinking_mode: str | None = None
+        self.thinking_budget: int | None = None
+        # Rolling conversation and action history for non-Markov policies
+        self.history_messages: list[dict[str, str]] = []  # chat-style without system
+        self.turn_index: int = 0
+        self.trajectory_history: list[dict[str, Any]] = []  # env/policy step records
+
+    async def initialize(self, config: dict[str, Any]) -> None:
+        # DEBUG: Log the incoming config
+        import logging
+        _logger = logging.getLogger(__name__)
+        _logger.debug(f"🔊 [POLICY_INIT] Received config keys: {list(config.keys())}")
+        _logger.debug(f"🔊 [POLICY_INIT] use_vision in config: {'use_vision' in config}, value: {config.get('use_vision')}")
+        _logger.debug(f"🔊 [POLICY_INIT] image_only_mode in config: {'image_only_mode' in config}, value: {config.get('image_only_mode')}")
+        
+        if "inference_url" in config:
+            self.inference_url = config["inference_url"]
+        if "model" in config:
+            self.model = config["model"]
+        if "use_tools" in config:
+            self.use_tools = bool(config["use_tools"])
+        if "use_vision" in config:
+            self.use_vision = bool(config["use_vision"])
+            _logger.debug(f"🔊 [POLICY_INIT] Set use_vision={self.use_vision} from config")
+        if "image_only_mode" in config:
+            self.image_only_mode = bool(config["image_only_mode"])
+            # If image_only_mode is enabled, automatically enable vision
+            if self.image_only_mode:
+                self.use_vision = True
+        # Auto-detect vision capability from model name if not explicitly set
+        if "use_vision" not in config and self.model:
+            self.use_vision = self._is_vision_model(self.model)
+        # Adopt sampling params from policy config (trainer passes these through)
+        if "temperature" in config:
+            self.temperature = float(config["temperature"])  # fail fast on bad types
+        if "top_p" in config:
+            self.top_p = float(config["top_p"])  # fail fast on bad types
+        if "max_tokens" in config:
+            self.max_tokens = int(config["max_tokens"])  # fail fast on bad types
+        # Thinking mode/budget forwarded into vLLM request (mirrors Wordle policy)
+        if "thinking_mode" in config:
+            self.thinking_mode = str(config["thinking_mode"])  # expect "think" or "no_think"
+        if "thinking_budget" in config and config["thinking_budget"] is not None:
+            self.thinking_budget = int(config["thinking_budget"])  # number of tokens inside <think>
+        if self.thinking_budget is None:
+            try:
+                if "openai.com" not in (self.inference_url or "").lower():
+                    self.thinking_budget = 1028
+            except Exception:
+                self.thinking_budget = 1028
+        # Reset state on (re)initialize
+        self.history_messages = []
+        self.turn_index = 0
+        self.trajectory_history = []
+        
+        # DEBUG: Log final state
+        _logger.debug(f"🔊 [POLICY_INIT] FINAL STATE: use_vision={self.use_vision}, image_only_mode={self.image_only_mode}, model={self.model}")
+
+    def _append_user_observation(self, observation_text: str) -> None:
+        self.history_messages.append({"role": "user", "content": observation_text})
+        self.turn_index += 1
+
+    def _append_assistant_turn(
+        self,
+        assistant_text: str | None,
+        tool_calls: list[dict[str, Any]] | None,
+        env_result: dict[str, Any] | None,
+    ) -> None:
+        # Record assistant content (if any)
+        if assistant_text is not None:
+            self.history_messages.append({"role": "assistant", "content": assistant_text})
+        # Keep structured step record for training/analysis
+        record: dict[str, Any] = {"turn": self.turn_index}
+        if tool_calls is not None:
+            record["tool_calls"] = tool_calls
+        if env_result is not None:
+            record["env_result"] = env_result
+        self.trajectory_history.append(record)
+
+    def build_inference_request(
+        self,
+        observation_text: str,
+        history: list[dict[str, Any]] | None = None,
+        turn: int | None = None,
+        image_parts: list[dict[str, Any]] | None = None,
+    ) -> dict[str, Any]:
+        messages = CrafterReActAgent.build_messages(
+            observation=observation_text,
+            history=history,
+            turn=turn,
+            image_parts=image_parts,
+            image_only_mode=self.image_only_mode,
+        )
+        
+        # DEBUG: Log message structure
+        import logging
+        _logger = logging.getLogger(__name__)
+        _logger.debug(f"🔊 [BUILD_REQUEST] Built {len(messages)} messages")
+        for idx, msg in enumerate(messages):
+            role = msg.get("role")
+            content = msg.get("content")
+            if isinstance(content, list):
+                _logger.debug(f"🔊 [BUILD_REQUEST] Message[{idx}] role={role}, content=list[{len(content)}]")
+                for part_idx, part in enumerate(content):
+                    if isinstance(part, dict):
+                        part_type = part.get("type")
+                        _logger.debug(f"🔊 [BUILD_REQUEST]   Part[{part_idx}]: type={part_type}")
+            else:
+                content_len = len(str(content)) if content else 0
+                _logger.debug(f"🔊 [BUILD_REQUEST] Message[{idx}] role={role}, content_len={content_len}")
+        
+        payload: dict[str, Any] = {
+            "messages": messages,
+        }
+        
+        # DEBUG: Verify messages are in payload correctly
+        _logger.debug(f"🔊 [BUILD_REQUEST_PAYLOAD] Created payload with {len(payload['messages'])} messages")
+        for idx, msg in enumerate(payload["messages"]):
+            content = msg.get("content")
+            _logger.debug(f"🔊 [BUILD_REQUEST_PAYLOAD] Payload message[{idx}]: type={type(content).__name__}, is_list={isinstance(content, list)}, len={len(content) if isinstance(content, list) else len(str(content)) if content else 0}")
+        
+        if self.model is not None:
+            payload["model"] = self.model
+        # Thinking controls
+        if self.thinking_mode is None and "openai.com" not in (self.inference_url or "").lower():
+            self.thinking_mode = "think"
+        if self.thinking_mode is not None:
+            payload["thinking_mode"] = self.thinking_mode
+        if self.thinking_budget is None and "openai.com" not in (self.inference_url or "").lower():
+            self.thinking_budget = 1028
+        if self.thinking_budget is not None:
+            payload["thinking_budget"] = self.thinking_budget
+        # Inject sampling parameters if set via initialize(config)
+        if self.max_tokens is not None:
+            # Use max_completion_tokens for newer models, max_tokens for older ones
+            if self.model and ("gpt-5" in self.model):
+                payload["max_completion_tokens"] = self.max_tokens
+            else:
+                payload["max_tokens"] = self.max_tokens
+        if self.temperature is not None:
+            payload["temperature"] = self.temperature
+        if self.top_p is not None:
+            payload["top_p"] = self.top_p
+        if self.use_tools:
+            payload["tools"] = TOOLS_SCHEMA
+            payload["tool_choice"] = "required"
+            payload["function_call"] = {"name": "interact_many"}
+            payload["parallel_tool_calls"] = False
+            # Ensure the inference server injects family-specific stop sequences
+            # to terminate immediately after the first tool call for compliance.
+            payload["stop_after_tool_calls"] = 1
+        return payload
+
+    @staticmethod
+    def parse_response_to_tool_calls(
+        response: dict[str, Any],
+        use_tools: bool = True,
+    ) -> list[dict[str, Any]]:
+        """Turn an inference response into environment tool calls."""
+        # First check if we got actual tool calls
+        choices = response.get("choices", [])
+        tool_calls: list[dict[str, Any]] = []
+
+        for choice in choices:
+            msg = choice.get("message", {})
+            if "tool_calls" in msg and msg["tool_calls"] is not None:
+                for tc in msg["tool_calls"]:
+                    if tc is None:
+                        continue
+                    # Handle both OpenAI format and simplified format
+                    if "function" in tc:
+                        # Standard OpenAI format
+                        tool_calls.append(
+                            {
+                                "tool_name": tc["function"]["name"],
+                                "arguments": tc["function"]["arguments"],
+                            }
+                        )
+                    elif "name" in tc:
+                        # Simplified format from our vLLM service
+                        tool_calls.append(
+                            {
+                                "tool_name": tc["name"],
+                                "arguments": tc["arguments"],
+                            }
+                        )
+
+        # If we got tool calls, return them
+        if tool_calls:
+            # Normalize common degenerate pattern ["move_right", "do"] when nothing is nearby.
+            # If previous env_result indicates no interaction target, drop trailing 'do'.
+            normalized: list[dict[str, Any]] = []
+            for tc in tool_calls:
+                if tc and isinstance(tc, dict) and tc.get("tool_name") == "interact_many":
+                    args = tc.get("arguments")
+                    if isinstance(args, str):
+                        try:
+                            import json
+
+                            args = json.loads(args)
+                        except (json.JSONDecodeError, ValueError):
+                            args = {}
+                    actions = []
+                    if isinstance(args, dict):
+                        maybe_actions = args.get("actions")
+                        if isinstance(maybe_actions, list):
+                            actions = maybe_actions
+                    # Simple heuristic: avoid repeating same pair; avoid 'do' with no context
+                    if len(actions) == 2 and actions[0] == "move_right" and actions[1] == "do":
+                        actions = ["move_right"]
+                    normalized.append(
+                        {"tool_name": "interact_many", "arguments": {"actions": actions or []}}
+                    )
+                else:
+                    normalized.append(tc)
+            return normalized
+
+        # No actions found
+        return []
+
+    async def step(
+        self,
+        observation_text: str,
+        state: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+        """Stateful step: update policy history and prepare inference request.
+
+        Inputs (via metadata, optional):
+        - "prev_assistant_text": str — assistant text from prior turn
+        - "prev_tool_calls": List[Dict] — tool calls executed last turn
+        - "prev_env_result": Dict — env step result for prior tool calls
+        - "prev_inference_response": Dict — raw LLM response; if present and
+          use_tools=False, we record assistant_text parsed from content.
+
+        Returns (tool_calls, meta):
+        - tool_calls: empty list; coordinator should call inference and then use
+          parse_response_to_tool_calls() to derive tool_calls
+        - meta: { inference_url, inference_request, turn_index, history_len }
+        """
+        # If caller provided results from previous cycle, record them first
+        if metadata is not None:
+            prev_assistant_text: str | None = None
+            prev_tool_calls: list[dict[str, Any]] | None = None
+            prev_env_result: dict[str, Any] | None = None
+            if "prev_assistant_text" in metadata:
+                prev_assistant_text = metadata["prev_assistant_text"]
+            if "prev_tool_calls" in metadata:
+                prev_tool_calls = metadata["prev_tool_calls"]
+            if "prev_env_result" in metadata:
+                prev_env_result = metadata["prev_env_result"]
+            if (
+                prev_assistant_text is not None
+                or prev_tool_calls is not None
+                or prev_env_result is not None
+            ):
+                self._append_assistant_turn(prev_assistant_text, prev_tool_calls, prev_env_result)
+
+        # Append current observation as the next user message (internal history only)
+        self._append_user_observation(observation_text)
+
+        # Build user message by combining the current observation text
+        # (formatted surroundings/inventory) with the previous 3 tool calls as context.
+        # Most recent first.
+        lines: list[str] = []
+
+        def _format_tool_call_line_for_context(
+            tool_name: str, arguments: Any, max_chars: int = 500
+        ) -> str:
+            import json as _json
+
+            # Render arguments compactly, then clip to max_chars
+            if isinstance(arguments, dict | list):
+                try:
+                    rendered = _json.dumps(arguments, ensure_ascii=False, separators=(",", ":"))
+                except Exception:
+                    rendered = str(arguments)
+            elif isinstance(arguments, str):
+                rendered = arguments
+            else:
+                rendered = str(arguments)
+            if isinstance(rendered, str) and len(rendered) > max_chars:
+                rendered = rendered[:max_chars]
+            return f"- {tool_name}: {rendered}"
+
+        # Prefer pulling from trajectory_history (accumulates over turns)
+        for record in reversed(self.trajectory_history):
+            if len(lines) >= 3:
+                break
+            tc_list = record.get("tool_calls")
+            if not tc_list:
+                continue
+            # Use the first tool call for that turn if multiple exist
+            tc = tc_list[0] if isinstance(tc_list, list) and tc_list else None
+            if not isinstance(tc, dict):
+                continue
+            name = tc.get("tool_name") or tc.get("name") or "unknown"
+            args = tc.get("arguments")
+            lines.append(_format_tool_call_line_for_context(name, args))
+
+        # If trajectory history is empty (first few turns), fall back to metadata once
+        if not lines and metadata is not None and metadata.get("prev_tool_calls"):
+            calls: list[dict[str, Any]] = metadata["prev_tool_calls"]
+            for call in reversed(calls):
+                if len(lines) >= 3:
+                    break
+                if not isinstance(call, dict):
+                    continue
+                name = call.get("tool_name") or call.get("name") or "unknown"
+                args = call.get("arguments")
+                lines.append(_format_tool_call_line_for_context(name, args))
+
+        context_text = "Previous tool calls (most recent first):\n" + (
+            "\n".join(lines) if lines else "- none"
+        )
+
+        # Combine observation with context so the model always sees surroundings/inventory
+        combined_text = f"{observation_text}\n\n{context_text}"
+
+        raw_observation: dict[str, Any] | None = None
+        if metadata is not None:
+            raw_candidate = metadata.get("raw_observation")
+            if isinstance(raw_candidate, dict):
+                raw_observation = raw_candidate
+        
+        # DEBUG: Log image extraction
+        import logging
+        _logger = logging.getLogger(__name__)
+        _logger.debug(f"🔊 [POLICY] use_vision={self.use_vision}, has_raw_obs={raw_observation is not None}")
+        if raw_observation:
+            obs = raw_observation.get("observation", raw_observation)
+            data_url = obs.get("observation_image_data_url") if isinstance(obs, dict) else None
+            _logger.debug(f"🔊 [POLICY] has_data_url={data_url is not None}, url_preview={data_url[:50] if data_url else 'NONE'}...")
+        
+        image_parts = self._extract_image_parts(raw_observation)
+        _logger.debug(f"🔊 [POLICY] Extracted {len(image_parts)} image parts")
+
+        payload = self.build_inference_request(
+            combined_text,
+            history=[],  # no prior user/assistant history
+            turn=self.turn_index,
+            image_parts=image_parts,
+        )
+        
+        # DEBUG: Verify payload before returning
+        _logger.debug(f"🔊 [POLICY_STEP_RETURN] About to return payload with {len(payload.get('messages', []))} messages")
+        for idx, msg in enumerate(payload.get("messages", [])):
+            content = msg.get("content")
+            _logger.debug(f"🔊 [POLICY_STEP_RETURN] Return message[{idx}]: type={type(content).__name__}, is_list={isinstance(content, list)}")
+            if isinstance(content, list):
+                _logger.debug(f"🔊 [POLICY_STEP_RETURN]   Content list has {len(content)} items")
+                # Add assertion to catch corruption early
+                assert len(content) > 0, f"Message content list is empty! This should contain images."
+        
+        meta_out = {
+            "inference_url": self.inference_url,
+            "inference_request": payload,
+            "turn_index": self.turn_index,
+            "history_len": len(self.history_messages),
+        }
+        return [], meta_out
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "turn_index": self.turn_index,
+            "history_messages": self.history_messages,
+            "trajectory_history": self.trajectory_history,
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.turn_index = int(state["turn_index"])
+        self.history_messages = state["history_messages"]
+        self.trajectory_history = state["trajectory_history"]
+
+    async def serialize(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "config": {
+                "inference_url": self.inference_url,
+                "model": self.model,
+                "use_tools": self.use_tools,
+                "use_vision": self.use_vision,
+            },
+            "state": self.state_dict(),
+        }
+
+    @classmethod
+    async def deserialize(cls, payload: dict[str, Any]) -> CrafterPolicy:
+        config = payload["config"]
+        state = payload["state"]
+        policy = cls(
+            inference_url=config["inference_url"],
+            model=config.get("model"),
+        )
+        policy.use_tools = bool(config.get("use_tools", True))
+        policy.use_vision = bool(config.get("use_vision", False))
+        policy.load_state_dict(state)
+        return policy
+
+    async def terminate(self) -> None:
+        return None
+
+    def prepare_inference_request(
+        self, observation: dict[str, Any], history: list[dict[str, Any]] = None
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]:
+        """Prepare an inference request (implementing abstract method)."""
+        # Format observation with rich contextual information
+        observation_text = self._format_observation_for_llm(observation)
+        
+        # Extract image parts based on vision settings
+        if self.use_vision:
+            image_parts = self._extract_image_parts(observation)
+        else:
+            # Text-only mode: don't include any images
+            image_parts = []
+
+        # Build messages with appropriate mode
+        messages = CrafterReActAgent.build_messages(
+            observation=observation_text,
+            history=history,
+            turn=self.turn_index,
+            image_parts=image_parts,
+            image_only_mode=self.image_only_mode,
+        )
+
+        # Return messages and tools schema
+        tools = TOOLS_SCHEMA if self.use_tools else None
+        return messages, tools
+
+    def _format_observation_for_llm(self, observation: dict[str, Any]) -> str:
+        """Format observation with rich contextual information for the LLM using the shared formatter."""
+        from .shared import format_observation
+
+        # Get the observation data (could be nested)
+        obs_data = observation.get("observation", observation)
+
+        # Ensure obs_data is a dict for safe access
+        if not isinstance(obs_data, dict):
+            return f"Observation: {str(observation)}"
+
+        # Use the shared format_observation function with step information
+        step_idx = observation.get("step_idx", 0)
+        max_steps = 100  # Default max steps, could be made configurable
+
+        # Get additional info from the observation wrapper
+        info = observation.get("info", {})
+        if isinstance(info, dict) and "health" in info and "health" not in obs_data:
+            obs_data = dict(obs_data)  # Make a copy
+            obs_data["health"] = info["health"]
+
+        return format_observation(obs_data, step_count=step_idx, max_steps=max_steps)
+
+    @staticmethod
+    def _is_vision_model(model_name: str) -> bool:
+        """Check if a model supports vision/image inputs based on its name."""
+        if not model_name:
+            return False
+        
+        model_lower = model_name.lower()
+        
+        # Known vision-capable model patterns
+        vision_patterns = [
+            "gpt-4o",           # GPT-4o series
+            "gpt-4-turbo",      # GPT-4 Turbo with vision
+            "gpt-4-vision",     # Explicit vision variant
+            "gpt-5",            # GPT-5 series (all variants support vision)
+            "claude-3",         # All Claude 3 models support vision
+            "gemini",           # Gemini models
+            "qwen-vl",          # Qwen Vision-Language models
+            "qwen3-vl",         # Qwen3 VL
+            "pixtral",          # Mistral's vision model
+            "llava",            # LLaVA models
+            "phi-3-vision",     # Microsoft Phi-3 Vision
+            "internvl",         # InternVL models
+            "cogvlm",           # CogVLM models
+            "vision",           # Generic vision indicator
+        ]
+        
+        return any(pattern in model_lower for pattern in vision_patterns)
+
+    def _extract_image_parts(
+        self, observation: dict[str, Any] | None
+    ) -> list[dict[str, Any]]:
+        """Extract image parts from crafter observation for vision-capable models.
+        
+        Returns OpenAI-style image_url format if vision is enabled and image data is available.
+        """
+        # Only extract images if vision is enabled for this policy
+        if not self.use_vision:
+            return []
+        
+        if not observation:
+            return []
+        
+        # Get the observation data (could be nested)
+        obs = observation.get("observation", observation)
+        if not isinstance(obs, dict):
+            return []
+        
+        # Extract the data URL (includes base64-encoded image)
+        data_url = obs.get("observation_image_data_url")
+        if not data_url or not isinstance(data_url, str):
+            return []
+        
+        # Return OpenAI-style image_url format
+        return [{"type": "image_url", "image_url": {"url": data_url}}]
+
+    def parse_model_response(
+        self, response: str, observation: dict[str, Any]
+    ) -> list[dict[str, Any]]:
+        """Parse model response into tool calls (implementing abstract method).
+
+        Note: Despite the type hint, vLLM actually returns a dict response,
+        not a string. We handle both cases.
+        """
+        # Handle dict response from vLLM (the actual case)
+        if isinstance(response, dict):
+            return self.parse_response_to_tool_calls(response, self.use_tools)
+
+        # Handle string response (fallback case for raw text)
+        if isinstance(response, str):
+            actions = CrafterReActAgent.parse_actions_from_response(response)
+            if actions:
+                return [{"tool_name": "interact_many", "arguments": {"actions": actions}}]
+
+        # Default empty response
+        return []
+
+
+__all__ = ["CrafterPolicy"]