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

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

@@ -1,28 +1,31 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Tuple
 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]], Optional[List[Dict[str, Any]]]]:
+        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]]:
+        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)
 
 
@@ -37,23 +40,23 @@ class CrafterPolicy(Policy):
 
     name: str = "crafter-react"
 
-    def __init__(self, inference_url: str, model: Optional[str] = None) -> None:
+    def __init__(self, inference_url: str, model: str | None = None) -> None:
         self.inference_url = inference_url
         self.model = model
         self.use_tools = True
         # Sampling parameters (populated via initialize(config))
-        self.temperature: Optional[float] = None
-        self.top_p: Optional[float] = None
-        self.max_tokens: Optional[int] = None
+        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: Optional[str] = None
-        self.thinking_budget: Optional[int] = None
+        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.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
+        self.trajectory_history: list[dict[str, Any]] = []  # env/policy step records
 
-    async def initialize(self, config: Dict[str, Any]) -> None:
+    async def initialize(self, config: dict[str, Any]) -> None:
         if "inference_url" in config:
             self.inference_url = config["inference_url"]
         if "model" in config:
@@ -89,15 +92,15 @@ class CrafterPolicy(Policy):
 
     def _append_assistant_turn(
         self,
-        assistant_text: Optional[str],
-        tool_calls: Optional[List[Dict[str, Any]]],
-        env_result: Optional[Dict[str, Any]],
+        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}
+        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:
@@ -107,13 +110,17 @@ class CrafterPolicy(Policy):
     def build_inference_request(
         self,
         observation_text: str,
-        history: Optional[List[Dict[str, str]]] = None,
-        turn: Optional[int] = None,
-    ) -> Dict[str, Any]:
+        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
+            observation=observation_text,
+            history=history,
+            turn=turn,
+            image_parts=image_parts,
         )
-        payload: Dict[str, Any] = {
+        payload: dict[str, Any] = {
             "messages": messages,
         }
         if self.model is not None:
@@ -148,9 +155,9 @@ class CrafterPolicy(Policy):
 
     @staticmethod
     def parse_response_to_tool_calls(
-        response: Dict[str, Any],
+        response: dict[str, Any],
         use_tools: bool = True,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Turn an inference response into environment tool calls.
 
         - If tools were used, expect tool_calls-compatible output and forward as-is
@@ -160,8 +167,8 @@ class CrafterPolicy(Policy):
         """
         # First check if we got actual tool calls
         choices = response.get("choices", [])
-        tool_calls: List[Dict[str, Any]] = []
-        
+        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:
@@ -185,18 +192,19 @@ class CrafterPolicy(Policy):
                                 "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]] = []
+            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 = {}
@@ -208,11 +216,13 @@ class CrafterPolicy(Policy):
                     # 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 []}})
+                    normalized.append(
+                        {"tool_name": "interact_many", "arguments": {"actions": actions or []}}
+                    )
                 else:
                     normalized.append(tc)
             return normalized
-            
+
         # Otherwise, parse plain text content for actions
         text = ""
         for choice in choices:
@@ -221,24 +231,25 @@ class CrafterPolicy(Policy):
             if content:
                 text = content
                 break
-        
+
         if text:
             # Try to parse actions from the text
             from .shared import parse_actions
+
             actions = parse_actions(text)
             if actions:
                 # Wrap actions in interact_many tool call
                 return [{"tool_name": "interact_many", "arguments": {"actions": actions}}]
-        
+
         # No actions found
         return []
 
     async def step(
         self,
         observation_text: str,
-        state: Optional[Dict[str, Any]] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        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):
@@ -255,16 +266,20 @@ class CrafterPolicy(Policy):
         """
         # If caller provided results from previous cycle, record them first
         if metadata is not None:
-            prev_assistant_text: Optional[str] = None
-            prev_tool_calls: Optional[List[Dict[str, Any]]] = None
-            prev_env_result: Optional[Dict[str, Any]] = 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:
+            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)
@@ -273,11 +288,15 @@ class CrafterPolicy(Policy):
         # 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:
+        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)):
+            if isinstance(arguments, dict | list):
                 try:
                     rendered = _json.dumps(arguments, ensure_ascii=False, separators=(",", ":"))
                 except Exception:
@@ -289,6 +308,7 @@ class CrafterPolicy(Policy):
             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:
@@ -306,7 +326,7 @@ class CrafterPolicy(Policy):
 
         # 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"]
+            calls: list[dict[str, Any]] = metadata["prev_tool_calls"]
             for call in reversed(calls):
                 if len(lines) >= 3:
                     break
@@ -316,17 +336,27 @@ class CrafterPolicy(Policy):
                 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")
+        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
+        image_parts = self._extract_image_parts(raw_observation)
+
         payload = self.build_inference_request(
             combined_text,
             history=[],  # no prior user/assistant history
             turn=self.turn_index,
+            image_parts=image_parts,
         )
-        #print("Debugging only:; ", payload)
+        # print("Debugging only:; ", payload)
         meta_out = {
             "inference_url": self.inference_url,
             "inference_request": payload,
@@ -335,19 +365,19 @@ class CrafterPolicy(Policy):
         }
         return [], meta_out
 
-    def state_dict(self) -> Dict[str, Any]:
+    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:
+    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]:
+    async def serialize(self) -> dict[str, Any]:
         return {
             "name": self.name,
             "config": {
@@ -359,7 +389,7 @@ class CrafterPolicy(Policy):
         }
 
     @classmethod
-    async def deserialize(cls, payload: Dict[str, Any]) -> "CrafterPolicy":
+    async def deserialize(cls, payload: dict[str, Any]) -> CrafterPolicy:
         config = payload["config"]
         state = payload["state"]
         policy = cls(
@@ -372,26 +402,28 @@ class CrafterPolicy(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]], Optional[List[Dict[str, Any]]]]:
+        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)
+        image_parts = self._extract_image_parts(observation)
 
         # Build messages (observation_text already formatted; no raw matrices)
         messages = CrafterReActAgent.build_messages(
             observation=observation_text,
             history=history,
-            turn=self.turn_index
+            turn=self.turn_index,
+            image_parts=image_parts,
         )
 
         # 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:
+    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
 
@@ -402,39 +434,43 @@ class CrafterPolicy(Policy):
         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):
-            # Merge health from info into obs_data for the formatter
-            if "health" in info and "health" not in obs_data:
-                obs_data = dict(obs_data)  # Make a copy
-                obs_data["health"] = info["health"]
+        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)
-    
+
+    def _extract_image_parts(
+        self, observation: dict[str, Any] | None
+    ) -> list[dict[str, Any]]:
+        """Crafter policy uses text-only prompts; do not attach image parts."""
+
+        return []
+
     def parse_model_response(
-        self, response: str, observation: Dict[str, Any]
-    ) -> List[Dict[str, Any]]:
+        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 []
 

examples/warming_up_to_rl/task_app/synth_envs_hosted/envs/crafter/react_agent.py

@@ -7,7 +7,7 @@ utilities to keep a single parser.
 
 from __future__ import annotations
 
-from typing import Dict, List, Optional
+from typing import Any
 
 from .shared import parse_actions
 
@@ -51,7 +51,7 @@ class CrafterReActAgent:
             "place_stone, place_table, place_furnace, place_plant, make_wood_pickaxe, make_stone_pickaxe, "
             "make_iron_pickaxe, make_wood_sword, make_stone_sword, make_iron_sword\n"
         )
-    
+
     @staticmethod
     def get_system_prompt_with_tools() -> str:
         """System prompt for tool-based interaction (e.g., Qwen3 models)."""
@@ -80,17 +80,29 @@ class CrafterReActAgent:
         )
 
     @staticmethod
-    def build_messages(observation: str, history: Optional[List[Dict[str, str]]] = None, turn: Optional[int] = None) -> List[Dict[str, str]]:
+    def build_messages(
+        observation: str,
+        history: list[dict[str, Any]] | None = None,
+        turn: int | None = None,
+        image_parts: list[dict[str, Any]] | None = None,
+    ) -> list[dict[str, Any]]:
         """Construct OpenAI-style messages list for vLLM generation."""
-        msgs: List[Dict[str, str]] = [{"role": "system", "content": CrafterReActAgent.get_system_prompt()}]
+        msgs: list[dict[str, Any]] = [
+            {"role": "system", "content": CrafterReActAgent.get_system_prompt()}
+        ]
         if history:
             msgs.extend(history)
-        msgs.append({"role": "user", "content": observation})
+        user_content: Any
+        if image_parts:
+            user_content = [{"type": "text", "text": observation}] + list(image_parts)
+        else:
+            user_content = observation
+        msgs.append({"role": "user", "content": user_content})
         return msgs
 
     @staticmethod
-    def parse_actions_from_response(response_text: str) -> List[str]:
+    def parse_actions_from_response(response_text: str) -> list[str]:
         return parse_actions(response_text)
 
 
-__all__ = ["CrafterReActAgent"]
+__all__ = ["CrafterReActAgent"]