My submission agent using just in time RL for action choosing

agent.py

@@ -26,12 +26,16 @@ Tips:
 import json
 import os
 import re
+import random
+import time
 from dataclasses import dataclass, field
 from typing import Optional
 
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
 
+from cross_episode_memory import CrossEpisodeMemory
+
 # Load environment variables
 load_dotenv()
 
@@ -67,7 +71,7 @@ else:
     LLM_CLIENT = InferenceClient(token=_hf_token)
 
 
-def call_llm(prompt: str, system_prompt: str, seed: int, max_tokens: int = 300) -> str:
+def call_llm(prompt: str, system_prompt: str, seed: int = 42, temperature: int = 0.0001, max_tokens: int = 1000) -> str:
     """
     Call the LLM with the given prompt. Use this function in your agent.
     
@@ -79,35 +83,32 @@ def call_llm(prompt: str, system_prompt: str, seed: int, max_tokens: int = 300)
         
     Returns:
         The LLM's response text
-        
-    Example:
-        response = call_llm(
-            prompt="You are in a forest. What do you do?",
-            system_prompt=SYSTEM_PROMPT,
-            seed=42,
-        )
     """
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt},
     ]
-
-    if USE_LOCAL_MODEL and _local_pipeline is not None:
-        outputs = _local_pipeline(
-            messages,
-            max_new_tokens=max_tokens,
-            temperature=0.0001,  # Near-deterministic (0.0 unsupported by some backends)
-            do_sample=True,
+    try:
+        if USE_LOCAL_MODEL and _local_pipeline is not None:
+            outputs = _local_pipeline(
+                messages,
+                max_new_tokens=max_tokens,
+                temperature=temperature,  # Near-deterministic (0.0 unsupported by some backends)
+                do_sample=True,
+            )
+            return outputs[0]["generated_text"][-1]["content"]
+
+        response = LLM_CLIENT.chat.completions.create(
+            model=LLM_MODEL,
+            messages=messages,
+            temperature=temperature,  # Deterministic for reproducibility
+            max_tokens=max_tokens,
+            seed=seed,
         )
-        return outputs[0]["generated_text"][-1]["content"]
-
-    response = LLM_CLIENT.chat.completions.create(
-        model=LLM_MODEL,
-        messages=messages,
-        temperature=0.0,  # Deterministic for reproducibility
-        max_tokens=max_tokens,
-        seed=seed,
-    )
+    
+    except Exception as e:
+        print(f"[LLM Error] {e}")
+        return "THOUGHT: LLM error, trying look."
 
     return response.choices[0].message.content
 
@@ -124,153 +125,746 @@ class RunResult:
     history: list[tuple[str, str, str]] = field(default_factory=list)
 
 
-# =============================================================================
-# System Prompt - Customize this for your agent
-# =============================================================================
-
-SYSTEM_PROMPT = """You are playing a classic text adventure game.
-
-GOAL: Explore the world, solve puzzles, and maximize your score.
-
-AVAILABLE TOOLS (use via MCP):
-- play_action: Execute a game command (north, take lamp, open mailbox, etc.)
-- memory: Get current game state and history (if implemented)
-- inventory: Check what you're carrying (if implemented)
-
-VALID GAME COMMANDS for play_action:
-- Movement: north, south, east, west, up, down, enter, exit
-- Objects: take <item>, drop <item>, open <thing>, close <thing>, examine <thing>
-- Other: look, inventory, read <thing>, turn on lamp
-
-RESPOND IN THIS EXACT FORMAT (no markdown):
-THOUGHT: <your reasoning about what to do next>
-TOOL: <tool_name>
-ARGS: <JSON arguments, e.g., {"action": "look"}>
-
-Example:
-THOUGHT: I should look around to see where I am.
-TOOL: play_action
-ARGS: {"action": "look"}
-"""
-
-
 # =============================================================================
 # Student Agent - IMPLEMENT THIS CLASS
 # =============================================================================
 
 class StudentAgent:
     """
-    Your ReAct agent implementation.
-    
-    TODO:
-    1. Implement the run() method with the ReAct loop
-    2. Parse LLM responses to extract tool calls
-    3. Track state and avoid loops
+    MCP ReAct Agent - A complete working example.
     
-    Use the provided call_llm() function to interact with the LLM.
+    This agent demonstrates:
+    - ReAct loop (Thought -> Tool -> Observation)
+    - Loop detection
+    - Action validation
+    - Score tracking via memory tool
     """
     
-    def __init__(self):
-        """Initialize your agent here."""
-        # TODO: Initialize any state tracking you need
-        # self.history = []
-        # self.visited_locations = set()
-        pass
+    def __init__(self, game: str = "zork1", guiding_prompt: str = None, top_actions: int = 5, exploration_alpha: float = 0.4, gamma: float = 0.95):
+        """Initialize the agent state."""
+        self.history: list[dict] = []
+        self.recent_actions: list[str] = []
+        self.score: int = 0
+        self.guiding_prompt = guiding_prompt or """Explore the environment and try to maximize your score."""
+        self.top_actions = top_actions
+
+        # Create game-specific memory directory
+        game_memory_dir = os.path.join("memory", game)
+        self.cross_mem = CrossEpisodeMemory(base_dir=game_memory_dir, eval_llm_model=LLM_MODEL, gamma=gamma)
+        self.exploration_alpha = exploration_alpha
     
     async def run(
         self,
-        client,  # FastMCP Client connected to your MCP server
+        client,
         game: str,
         max_steps: int,
         seed: int,
         verbose: bool = False,
     ) -> RunResult:
-        """
-        Run the agent for a game session.
+        """Run the agent for a game session."""
+        locations_visited = set()
+        history = []
+        moves = 0
         
-        Args:
-            client: FastMCP Client connected to your MCP server
-            game: Name of the game being played (e.g., "zork1")
-            max_steps: Maximum number of steps to take
-            seed: Random seed for reproducibility (use for LLM calls)
-            verbose: Whether to print detailed output
-            
-        Returns:
-            RunResult with final score and statistics
-        """
-        # TODO: Implement your ReAct loop here
-        #
-        # Basic structure:
-        # 1. Get initial observation (call play_action with "look")
-        # 2. Loop for max_steps:
-        #    a. Build prompt with current observation and history
-        #    b. Call LLM to get thought and action
-        #    c. Parse the response to extract tool and args
-        #    d. Call the tool via client.call_tool(tool_name, args)
-        #    e. Update history and state
-        #    f. Check for game over
-        # 3. Return RunResult with final statistics
+        # Get list of available tools
+        tools = await client.list_tools()
+        tool_names = [t.name for t in tools]
         
-        # Example of calling a tool:
-        # result = await client.call_tool("play_action", {"action": "look"})
-        # observation = result[0].text if result else "No response"
+        # Get initial observation
+        result = await client.call_tool("reset_game", {})
+        raw = self._extract_result(result)
+
+        payload = json.loads(raw)
+        state = payload["state"]          # <-- dict
+        observation = state["observation"]
         
-        # Example of calling the LLM:
-        # response = call_llm(
-        #     prompt="Current observation: " + observation,
-        #     system_prompt=SYSTEM_PROMPT,
-        #     seed=seed,
-        # )
+        # Track initial location
+        location = observation.split("\n")[0] if observation else "Unknown"
+        locations_visited.add(location)
         
-        # Placeholder implementation - replace with your code
-        locations_visited = set()
-        history = []
-        final_score = 0
-        moves = 0
+        if verbose:
+            print(f"\n{observation}")
         
-        # TODO: Your implementation here
-        # ...
+        # Parse initial score
+        self._update_score(state)
+        
+        # Main ReAct loop
+        for step in range(1, max_steps + 1):
+            # Build prompt with context
+            user_prompt, sys_prompt, memory_text = self._build_prompt(state)
+            
+            # Get response format 
+            response_format = self._get_response_format()
+            
+            # Call LLM for reasoning (use step-based seed for variety)
+            response = call_llm(user_prompt, sys_prompt, seed=seed + step, response_format=response_format)
+            
+            # Parse the response
+            action, options_with_confidences = self._parse_response(response)
+            
+            # Update action scores with memory
+            updated_options_with_logits = self.update_scores_with_memory(
+                observation, 
+                options_with_confidences,
+                k=10,
+                r=0.9,
+                memory_text=memory_text)
+            
+            if updated_options_with_logits:
+                # Select action with highest corrected confidence
+                valid_actions = state.get("valid_actions", []) or []
+                updated_options_valid_items = []
+                for option in updated_options_with_logits.items():
+                    if option[1]['action'] in valid_actions:
+                        updated_options_valid_items.append(option)
+                
+                best_option = max(
+                    updated_options_valid_items,
+                    key=lambda x: x[1].get('corrected_confidence', 0)
+                )
+                action = best_option[1]['action']
+            
+            if verbose:
+                print(f"\n--- Step {step} | Score: {self.score} ---")
+                print(f"[RAW_LLM_OUTPUT]\n{response}")
+                print(f"[ACTION] {action}")
+            
+            # Loop detection
+            self.recent_actions.append(action)
+            if len(self.recent_actions) > 5:
+                self.recent_actions = self.recent_actions[-5:]
+                
+            # Detect loops - if same action 3 times, force "look"
+            if len(self.recent_actions) >= 3 and len(set(self.recent_actions[-3:])) == 1:
+                if verbose:
+                    print(f"[WARNING] Loop detected - forcing 'look'")
+                action = "look"
+                self.recent_actions.append(action)
+                
+            moves += 1
+
+            # Update history
+            self.history.append({
+                "old_state": observation,
+                "action": action,
+                "full_response": response,
+                "reward": None,
+                "score": None,
+            })
+            
+            # Execute the tool
+            try:
+                result = await client.call_tool("play_action", {"action": action})
+                raw = self._extract_result(result)
+                payload = json.loads(raw)
+                state = payload["state"]
+                observation = state["observation"]
+                
+                if verbose:
+                    print(f"[RESULT] {observation[:200]}...")
+            except Exception as e:
+                observation = f"Error: {e}"
+                if verbose:
+                    print(f"[ERROR] {e}")
+            
+            # Track location
+            location = observation.split("\n")[0] if observation else "Unknown"
+            locations_visited.add(location)
+            
+            old_score = self.score
+            
+            # Track score from observation
+            self._update_score(state)
+            
+            self.history[-1]["reward"] = self.score - old_score
+            self.history[-1]["score"] = self.score
+            self.history[-1]["state"] = observation
+            
+            # Record in result history
+            history.append((response, f"play_action(action: {action})", observation[:100]))
+            
+            # Check for game over
+            if self._is_game_over(observation):
+                if verbose:
+                    print("\n*** GAME OVER ***")
+                break
+        
+        success = self._is_game_over(observation) and self.score > 0
+        
+        # Adding the current execution to memory for future episodes
+        self.cross_mem.add_episode(
+            game_history=self.history,
+            final_score=self.score,
+            success=success,
+            state=observation
+        )
         
         return RunResult(
-            final_score=final_score,
-            max_score=350,  # Zork1 max score, adjust if needed
+            final_score=self.score,
+            max_score=350,
             moves=moves,
             locations_visited=locations_visited,
-            game_completed=False,
+            game_completed=self._is_game_over(observation),
             history=history,
         )
     
-    def _build_prompt(self, observation: str, history: list) -> str:
-        """
-        Build the prompt for the LLM.
+    def update_scores_with_memory(self, current_state, options_with_confidences, k, r, memory_text, exploration_prob=0.4):
+        nearest_trajectories = self.cross_mem.retrieve_similar(
+                game_history=self.history,
+                current_state=current_state,
+                current_summary=memory_text,
+                k=k,
+                r=r
+            )
+        if not nearest_trajectories:
+            return {}
         
-        TODO: Implement this to create effective prompts
-        """
-        # TODO: Combine system prompt, history, and current observation
-        pass
-    
-    def _parse_response(self, response: str) -> tuple[str, str, dict]:
-        """
-        Parse LLM response to extract thought, tool name, and arguments.
+        # Get existing actions
+        existing_actions = set()
+        for option_data in options_with_confidences.values():
+            if isinstance(option_data, dict) and 'action' in option_data:
+                existing_actions.add(option_data['action'])
+
+        # Aggregate action rewards from nearest_trajectories
+        action_rewards = {}
+        for result_dict in nearest_trajectories:
+            action = result_dict.get('action', '').strip()
+            discounted_reward = result_dict.get('discounted_reward', 0)
+            if action not in action_rewards:
+                action_rewards[action] = []
+            action_rewards[action].append(discounted_reward)
         
-        TODO: Implement robust parsing
+        # Look for the action in the options suggested by the LLM
+        for action in action_rewards:
+            found = False
+            for option_data in options_with_confidences.values():
+                if isinstance(option_data, dict) and option_data.get('action', '') == action:
+                    found = True
+                    break
+            if not found:
+                rewards = action_rewards[action]
+                if len(rewards) >= 1 and sum(rewards) / len(rewards) > 0:
+                    if options_with_confidences:
+                        new_option_num = max(options_with_confidences.keys()) + 1
+                    else:
+                        new_option_num = 1
+                    options_with_confidences[new_option_num] = {
+                        'action': action,
+                        'confidence': 0
+                    }
         
+        # Calculate average discounted reward for each action Q(s, a)
+        action_avg_rewards = {}
+        for action, rewards in action_rewards.items():
+            if rewards:  # Ensure not empty
+                action_avg_rewards[action] = sum(rewards) / len(rewards)
+            else:
+                action_avg_rewards[action] = 0
+        
+        # Calculate overall average discounted reward (baseline) "state value"
+        count = 0
+        if action_rewards:
+            all_rewards = []
+            for rewards_list in action_rewards.values():
+                all_rewards.extend(rewards_list)
+            count = len(all_rewards)
+            overall_avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0
+        else:
+            overall_avg_reward = 0
+            
+        # The second part of step 4 (The exploration bonus) : Completing the Q(s, a) for the actions that are not in memory
+        alpha = self.exploration_alpha
+        new_actions_added = False
+        for option_num, option_data in options_with_confidences.items():
+            action = option_data['action']
+            if action not in action_avg_rewards:
+                rad = random.random()
+                print(f"action: {action} rad:{rad}")
+                if rad < exploration_prob:
+                    # UCB-style exploration bonus: decreases as sample count increases
+                    if count > 0:
+                        action_avg_rewards[action] = overall_avg_reward + alpha / count
+                    else:
+                        action_avg_rewards[action] = overall_avg_reward + alpha
+                else:
+                    action_avg_rewards[action] = 0
+                new_actions_added = True
+        
+        
+        # If new actions were added, recalculate overall_avg_reward
+        if new_actions_added:
+            all_avg_rewards = list(action_avg_rewards.values())
+            overall_avg_reward = sum(all_avg_rewards) / len(all_avg_rewards) if all_avg_rewards else 0
+
+        # Calculate advantage value for each action
+        action_advantages = {}
+        for action, avg_reward in action_avg_rewards.items():
+            action_advantages[action] = avg_reward - overall_avg_reward
+        
+        # Normalize advantage values
+        if action_advantages:
+            adv_values = list(action_advantages.values())
+            
+            positive_advs = [adv for adv in adv_values if adv > 0]
+
+            if positive_advs:
+                max_positive = max(positive_advs)
+                normalized_advantages = {action: adv / max_positive for action, adv in action_advantages.items()}
+            else:
+                max_negative_abs = abs(min(adv_values))
+                if max_negative_abs > 0:
+                    normalized_advantages = {action: adv / max_negative_abs for action, adv in action_advantages.items()}
+                else:
+                    normalized_advantages = {action: 0 for action in action_advantages.keys()}
+        else:
+            normalized_advantages = {}
+        
+        updated_options = {}
+        if options_with_confidences and normalized_advantages:
+            print(f"\n=== Action Advantage Analysis & Logit Correction ===")
+            print(f"Overall action average reward baseline: {overall_avg_reward:.4f}")
+            
+            for option_num, option_data in options_with_confidences.items():
+                if isinstance(option_data, dict) and 'action' in option_data:
+                    action = option_data['action']
+                    
+                    # Get normalized advantage value
+                    normalized_advantage = normalized_advantages.get(action, 0)
+                    raw_advantage = action_advantages.get(action, 0)
+                    avg_reward = action_avg_rewards.get(action, 0)
+
+                    # Calculate episode-based weight for normalized advantage
+                    # Weight increases from 1.0 (episode 1) to 1.5 (episode 50)
+                    # Formula: weight = 1.0 + (current_episode / 50) * 0.5
+                    # Clamped to max of 1.5 for episodes beyond 50
+
+                    if self.cross_mem:
+                        current_episode = self.cross_mem.current_episode_number
+                        episode_weight = min(1.0 + (current_episode / 50.0) * 0.5, 1.5)
+                    else:
+                        episode_weight = 1.0
+
+                    # Apply episode weight to normalized advantage
+                    weighted_normalized_advantage = normalized_advantage * episode_weight
+
+                    # Correct logit (add weighted normalized advantage value)
+                    normalized_prob = option_data.get('confidence', 0)
+                    corrected_logprob = normalized_prob + weighted_normalized_advantage
+                    
+                    # Create corrected option data
+                    updated_options[option_num] = {
+                        'action': action,
+                        'normalized_advantage': normalized_advantage,
+                        'corrected_confidence': corrected_logprob,
+                        'avg_reward': avg_reward,
+                        'raw_advantage': raw_advantage,
+                        'confidence': option_data.get('confidence', 0)
+                    }
+                    
+                    print(f"  Action: {action}")
+                    print(f"    Average reward: {avg_reward:.4f} | Raw advantage: {raw_advantage:.4f}")
+                    print(f"    Normalized advantage: {normalized_advantage:.4f}")
+                    print(f"    Original logprob: {normalized_prob:.4f} -> Corrected: {corrected_logprob:.4f}")
+                else:
+                    # If no matching advantage value, keep original value
+                    updated_options[option_num] = option_data.copy()
+                    updated_options[option_num]['corrected_confidence'] = option_data.get('confidence', 0)
+                    updated_options[option_num]['normalized_advantage'] = 0
+            print("="*50)
+            return updated_options
+        
+        else:
+            return {}
+    
+    def _build_prompt(self, state: dict) -> tuple[str, str]:
+        """Build the prompt for the LLM with context."""
+
+        valid_actions = state.get("valid_actions", []) or [] 
+        current_inventory_list = state.get("inventory", []) or []
+        current_inventory = ", ".join(current_inventory_list) if current_inventory_list  else "(empty)"
+        
+        if self.history:
+            summary, memory_text = self.generate_history_summary(self.history, state.get("observation", ""), temperature=0.0, max_tokens=1000, current_inventory=current_inventory)
+        else:
+            summary = ""
+            memory_text = "No game history."
+            # Add inventory even if no game history
+            if current_inventory:
+                memory_text += f"\n\n[INVENTORY]\n{current_inventory}"
+        sys_prompt = """You are an expert player aiming to complete a text-based adventure game. Points are given for making progress in the game. Select promising actions based on the game state and memory of past interactions.
+
+**EXPLORATION PRIORITY**: When you arrive at a NEW location you haven't fully explored before, you MUST thoroughly explore it FIRST before leaving."""
+        if self.guiding_prompt:
+            sys_prompt += f"\n\nFollow this guide: {self.guiding_prompt}"
+
+        # Add reminder about action list order only when valid_actions are provided
+        if valid_actions:
+            sys_prompt += """\n\n**CRITICAL CONSTRAINT**: When REFERENCE ACTIONS are provided, you MUST ONLY choose actions from that list. Any action not in the REFERENCE ACTIONS list is INVALID and will fail. Do NOT create custom actions. The list is unordered - position doesn't indicate quality."""
+
+        recent_history = ""
+        if self.history:
+            # Only keep the last 10 steps
+            recent_game_history = self.history[-5:]
+            start_index = max(0, len(self.history)-5)
+            for idx, entry in enumerate(recent_game_history):
+                actual_step = start_index + idx
+                recent_history += f"Step {actual_step}:\n"
+                recent_history += f"State: {entry.get('state', '')}\n"
+                recent_history += f"Action: {entry.get('action', '')}\n"
+                if entry.get('reward') is not None:
+                    recent_history += f"Reward: {entry.get('reward', 0)}\n"
+                recent_history += "\n"
+
+        # Add valid actions section if available
+        valid_actions_text = ""
+        if valid_actions:
+            valid_actions = valid_actions[:]  # make a copy to avoid side effects
+            rng = random.Random(time.time_ns())
+            rng.shuffle(valid_actions)
+            print(f"Valid actions (shuffled): {valid_actions}")
+            valid_actions_text = f"\nREFERENCE ACTIONS (ONLY VALID ACTIONS):\n{valid_actions}\n\n**STRICT REQUIREMENT**: These are the ONLY valid actions for the current state. You MUST select your actions EXCLUSIVELY from this list. Any action not in this list is INVALID and will be rejected by the game. Do NOT create, modify, or suggest any custom actions.**\n"
+
+        # Build option fields with confidence
+        option_fields = []
+        for i in range(1, self.top_actions + 1):
+            option_fields.append(f'    "reasoning{i}": "Why this action makes sense",')
+            option_fields.append(f'    "option{i}": "action command",')
+            option_fields.append(f'    "confidence{i}": 80,')
+        option_format = '\n'.join(option_fields)
+
+        user_prompt = f"""
+GAME HISTORY:
+{summary}
+{memory_text}
+
+RECENT STEPS:
+{recent_history}
+
+CURRENT STATE: {state.get("observation", "")}
+
+TASK:
+1. Analyze your progress: What have you achieved? What's your next objective?
+2. **MANDATORY: Check the REFERENCE ACTIONS list below - you MUST ONLY select from this list**
+3. Propose {self.top_actions} different actions with reasoning
+4. For EACH action, provide your confidence as an integer from 0 to 100 (e.g., 80 means 80% confidence that this action will help achieve the goal)
+
+RESPONSE FORMAT (JSON):
+{{
+    "progress_analysis": "What you've achieved and current challenges",
+    "next_objective": "Your next goal",
+{option_format}
+    "best_action": 1
+}}
+
+IMPORTANT:
+- **ABSOLUTE REQUIREMENT**: ALL actions (option1, option2, etc.) MUST be selected EXACTLY from the REFERENCE ACTIONS list below
+- Actions NOT in the REFERENCE ACTIONS list are INVALID and will cause the game to fail
+- DO NOT create custom actions, DO NOT modify actions from the list, DO NOT combine actions
+- **CRITICAL**: The sum of all confidence values (confidence1 + confidence2 + ... + confidence{self.top_actions}) MUST equal 100
+- The confidence values MUST have meaningful differences between them
+- Higher confidence means you believe this action is more likely to succeed
+- Lower confidence means you believe this action is less likely to succeed
+- Pay attention to game hints and clues in state descriptions
+- Don't repeat failed actions or create loops (e.g., north→south→north)
+
+{valid_actions_text}
+"""
+        return user_prompt, sys_prompt, memory_text
+    
+    def generate_history_summary(self, game_history, current_state=None, temperature=0.8, max_tokens=1000, current_inventory=None):
+        """
+        Generate structured LLM-summarized history context.
+        Summarizes the past states and optionally includes the current state.
+
+        Args:
+            game_history: List of game history entries containing state, action, score, reward, etc.
+            current_state: The current state to include in the summary (optional)
+            llm_model: The LLM model to use for generation
+            temperature: Temperature for LLM generation
+            max_tokens: Maximum tokens for the summary
+            current_inventory: The current inventory to include in the summary (optional)
+
         Returns:
-            Tuple of (thought, tool_name, args_dict)
+            tuple: (summary_text, structured_summary)
+                - summary_text: Natural language summary under [SUMMARY]
+                - structured_summary: Structured sections [PROGRESS], [LOCATION], [NEXT_OBJECTIVE], [INVENTORY]
         """
-        # TODO: Parse the response format:
-        # THOUGHT: ...
-        # TOOL: ...
-        # ARGS: {...}
-        pass
+        if not game_history and not current_state:
+            return "", ""
+
+        # Extract states and actions from game_history
+        states = [entry.get('state', '') for entry in game_history]
+        actions = [entry.get('action', '') for entry in game_history]
+        current_step = len(game_history)  # Current step index
+
+        # Build trajectory text for all past steps
+        earlier_trajectory_text = ""
+        for i in range(current_step):
+            earlier_trajectory_text += f"Step {i}:\n"
+            if i < len(states):
+                earlier_trajectory_text += f"State: {states[i]}\n"
+            if i < len(actions) and actions[i]:
+                earlier_trajectory_text += f"Action: {actions[i]}\n"
+            earlier_trajectory_text += "\n"
+
+        # Add current state if provided
+        if current_state:
+            earlier_trajectory_text += f"Step {current_step} (Current):\n"
+            earlier_trajectory_text += f"State: {current_state}\n"
+
+        sys_prompt = """You are an expert at analyzing game trajectories and creating highly distinctive summaries.
+Your milestones must be CONCISE and DISTINCTIVE - use specific keywords that clearly differentiate different game states.
+
+Output Format Requirements:
+[SUMMARY]: Provide a natural language summary that describes:
+   - The game's objective or goal (inferred from the trajectory)
+   - Current progress toward that goal
+   - Key accomplishments so far
+   - Current situation/status
+
+[PROGRESS]: List milestones in format "✓ M#: <action>→<key object/result>"
+   - If NO steps have score increases, output ONLY "No Progress" (no milestones)
+
+[LOCATION]: Track location changes in format "Location: A→B→C"
+   - Extract location names from state descriptions
+   - Only record when location actually changes
+   - CRITICAL: IGNORE unproductive loops - if the player goes back and forth between locations WITHOUT score increases, skip those redundant location movements
+   - Focus on the MEANINGFUL location trajectory that led to progress or new discoveries
+   - Example: If player went "room A→room B→room A→room B→room A→room C" with no score increase for the A↔B movements, record as "room A→room C"
+
+Milestone Format Rules:
+- Use format: <verb>→<critical object/state>
+- Focus on STATE CHANGES, not descriptions
+- Use specific nouns not generic terms
+- Each milestone should capture ONE concrete action or discovery
+- CRITICAL: Only record steps where the score increased compared to the previous step
+
+Examples:
+BAD SUMMARY: "The player has been exploring."
+GOOD SUMMARY: "The game's objective is to escape the haunted mansion. So far, the player has found a key in the library and unlocked the basement door, earning 15 points. Currently in the dark basement, the player needs to find a light source to continue exploring."
+
+BAD:  "✓ Milestone 1: Entered the library and proceeded to the ground floor stacks"
+GOOD: "✓ M1: enter→library ground floor"
+
+Location Examples:
+BAD:  "Location: You are in a library→You moved to another area→You are somewhere else"
+GOOD: "Location: entrance→library→north corridor→atrium"
+BAD:  "Location: room A→room B→room A→room B→room A" (unproductive loop)
+GOOD: "Location: room A" (ignore the meaningless back-and-forth)
+"""
+
+        user_prompt = f"""Analyze the game trajectory and generate a natural language summary followed by structured sections.
+
+Game History:
+{earlier_trajectory_text}
+
+Generate summary in this format:
+
+[SUMMARY]
+<Natural language paragraph (2-4 sentences) describing: game objective, current progress, key accomplishments, current situation>
+
+[PROGRESS]
+✓ M1: <verb>→<result>
+✓ M2: <verb>→<result>
+...
+
+[LOCATION]
+Location: A→B→C→...
+
+Requirements:
+- SUMMARY: Write a natural language paragraph explaining the game's goal and current status
+- Use specific object names
+- Focus on concrete state changes and discoveries
+- Avoid repetitive phrasing between milestones
+- CRITICAL: ONLY record progress where the score increased from the previous step
+- IMPORTANT: If there are NO steps with score increases, output ONLY "No Progress" under [PROGRESS] section (do not create any milestone items)
+- For locations: Extract actual location names from state descriptions and track when they change
+- Location format must be: "Location: " followed by location names separated by →
+- CRITICAL FOR LOCATION: Detect and REMOVE unproductive loops from the location trajectory
+  * Look at the score/reward at each step
+  * If the player moves back and forth between locations WITHOUT any score increase, those movements are loops - SKIP them
+  * Only include location changes that either: (1) led to score increases, OR (2) represent meaningful forward exploration
+  * Example: If steps show "A→B (no reward)→A (no reward)→B (no reward)→A (no reward)→C (reward +5)", output "Location: A→C"
+
+Generate the complete summary:"""
+
+        full_response = call_llm(user_prompt, sys_prompt, temperature=temperature, max_tokens=max_tokens).strip()
+
+        if full_response == "":
+            print("[Warning] No response from LLM for history summary generation")
+
+        # Extract [SUMMARY] section
+        summary_text = ""
+        if "[SUMMARY]" in full_response:
+            summary_start = full_response.find("[SUMMARY]") + len("[SUMMARY]")
+            # Find the next section marker
+            next_section = full_response.find("[PROGRESS]", summary_start)
+            if next_section == -1:
+                next_section = full_response.find("[LOCATION]", summary_start)
+
+            if next_section != -1:
+                summary_text = full_response[summary_start:next_section].strip()
+            else:
+                summary_text = full_response[summary_start:].strip()
+
+        # Extract structured sections (everything except [SUMMARY])
+        structured_summary = full_response
+        if "[SUMMARY]" in full_response:
+            # Remove [SUMMARY] section from structured summary
+            progress_start = full_response.find("[PROGRESS]")
+            if progress_start != -1:
+                structured_summary = full_response[progress_start:]
+
+        # Add inventory section to structured summary if available
+        if current_inventory:
+            structured_summary += f"\n\n[INVENTORY]\n{current_inventory}"
+
+        return summary_text, structured_summary
+
+    def _get_response_format(self) -> dict:
+        # Dynamically create JSON Schema to force final choice to be valid option numbers only
+        properties = {}
+        required_fields = []
+        valid_choices = []
+
+        # Add progress analysis fields
+        properties["progress_analysis"] = {
+            "type": "string",
+            "description": "Analysis of achievements and progress so far"
+        }
+        properties["next_objective"] = {
+            "type": "string",
+            "description": "Overall objective for the next steps"
+        }
+        required_fields.extend(["progress_analysis", "next_objective"])
+
+        # Dynamically generate reasoning and option fields
+        for i in range(1, self.top_actions + 1):
+            reasoning_key = f"reasoning{i}"
+            option_key = f"option{i}"
+
+            properties[reasoning_key] = {
+                "type": "string",
+                "description": f"Reasoning for option {i}"
+            }
+            properties[option_key] = {
+                "type": "string",
+                "description": f"Option {i} possible action"
+            }
+            required_fields.append(reasoning_key)
+            required_fields.append(option_key)
+            valid_choices.append(i)
+
+            confidence_key = f"confidence{i}"
+            properties[confidence_key] = {
+                "type": "integer",
+                "minimum": 0,
+                "maximum": 100,
+                "description": f"Confidence for option {i} (0-100)"
+            }
+            required_fields.append(confidence_key)
+
+        # Add best_action field
+        properties["best_action"] = {
+            "type": "number",
+            "minimum": 1,
+            "maximum": self.top_actions,
+            "description": f"The number of the best option (must be one of: {valid_choices})"
+        }
+        required_fields.extend([ "best_action"])
+        
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "game_action_choice",
+                "strict": True,
+                "schema": {
+                    "type": "object",
+                    "properties": properties,
+                    "required": required_fields,
+                    "additionalProperties": False
+                }
+            }
+        }
+
+        return response_format
     
-    def _call_llm(self, prompt: str, system_prompt: str, seed: int) -> str:
-        """
-        Call the LLM with the given prompt.
+    def _parse_response(self, response: str) -> tuple[str, dict]:
+        """Parse the LLM response to extract thought, tool, and arguments."""
+        # Parse JSON response
+        # Clean up markdown code blocks if present
+        cleaned_response = response.strip()
+        if cleaned_response.startswith('```'):
+            # Remove ```json or ``` at the start and ``` at the end
+            cleaned_response = re.sub(r'^```(?:json)?\s*\n?', '', cleaned_response)
+            cleaned_response = re.sub(r'\n?```\s*$', '', cleaned_response)
+
+        try:
+            json_response = json.loads(cleaned_response)
+        except json.JSONDecodeError as e:
+            print(f"JSON parsing error: {e}")
+            print(f"Raw response: {response}")
+            return "look", {}
+
+        # Extract and print progress analysis
+        progress_analysis = json_response.get("progress_analysis", "No analysis provided")
+        next_objective = json_response.get("next_objective", "No objective specified")
+
+        print("\n" + "="*50)
+        print("=== PROGRESS ANALYSIS ===")
+        print("="*50)
+        print(f"Achievements: {progress_analysis}")
+        print(f"\nNext Objective: {next_objective}")
+        print("="*50 + "\n")
+
+        # Extract all options with their reasoning and confidence
+        options_with_confidences = {}
         
-        This is a convenience wrapper - you can also use call_llm() directly.
-        """
-        return call_llm(prompt, system_prompt, seed)
+        for i in range(1, self.top_actions + 1):
+            reasoning_key = f"reasoning{i}"
+            option_key = f"option{i}"
+            confidence_key = f"confidence{i}"
+            
+            reasoning = json_response.get(reasoning_key, "No reasoning provided")
+            option = json_response.get(option_key, "look")
+            confidence = json_response.get(confidence_key, 50)
+            
+            options_with_confidences[i] = {
+                'action': option.strip(),
+                'reasoning': reasoning,
+                'confidence': confidence / 100.0
+            }
+
+        # Get the best action
+        best_choice = json_response.get("best_action", 1)
+        action_text = options_with_confidences.get(best_choice, {}).get('action', "look")
+
+        return action_text, options_with_confidences
+    
+    def _extract_result(self, result) -> str:
+        """Extract text from MCP tool result."""
+        if hasattr(result, 'content') and result.content:
+            return result.content[0].text
+        if isinstance(result, list) and result:
+            return result[0].text if hasattr(result[0], 'text') else str(result[0])
+        return str(result)
+    
+    def _update_score(self, state: dict) -> None:
+        if not state:
+            return
+        if "score" in state:
+            self.score = int(state["score"])
+    
+    def _is_game_over(self, text: str) -> bool:
+        """Check if the game is over."""
+        game_over_phrases = [
+            "game over",
+            "you have died",
+            "you are dead",
+            "*** you have died ***",
+        ]
+        text_lower = text.lower()
+        return any(phrase in text_lower for phrase in game_over_phrases)
 
 
 # =============================================================================

cross_episode_memory.py

@@ -0,0 +1,442 @@
+"""
+Cross-Episode Memory for Jericho Agent with Vector Similarity Search
+Uses state embeddings to find similar past experiences and update action selection.
+
+Installation:
+    pip install sentence-transformers faiss-cpu numpy --break-system-packages
+
+Alternative models:
+    - 'all-MiniLM-L6-v2' (default, 384 dims, fast)
+    - 'all-mpnet-base-v2' (768 dims, better quality)
+    - 'paraphrase-multilingual-MiniLM-L12-v2' (multilingual)
+"""
+
+import os
+import json
+import pickle
+from typing import List, Dict, Any, Optional, Tuple
+from collections import defaultdict, Counter
+import numpy as np
+from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import InferenceClient
+import faiss
+import traceback
+from openai import OpenAI
+
+load_dotenv()
+
+# Set USE_LOCAL_MODEL=1 in your .env to use a locally downloaded model
+USE_LOCAL_MODEL = os.getenv("USE_LOCAL_MODEL", "0").strip() in ("1", "true", "yes")
+LOCAL_MODEL_ID = os.getenv("LOCAL_MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")
+
+# Initialize the LLM client based on mode
+_local_pipeline = None
+
+if USE_LOCAL_MODEL:
+    import torch
+    from transformers import pipeline as _hf_pipeline
+
+    _local_pipeline = _hf_pipeline(
+        "text-generation",
+        model=LOCAL_MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+    LLM_CLIENT = None
+else:
+    _hf_token = os.getenv("HF_TOKEN")
+    if not _hf_token:
+        raise ValueError("HF_TOKEN not found. Set it in your .env file.")
+    LLM_CLIENT = InferenceClient(token=_hf_token)
+
+class CrossEpisodeMemory:    
+    def __init__(
+        self,
+        base_dir: str,
+        eval_llm_model: str = "Qwen/Qwen2.5-72B-Instruct",
+        gamma: float = 0.95,
+        embedding_model: str = 'all-MiniLM-L6-v2',
+        vector_dim: int = 384
+    ):
+        self.base_dir = base_dir
+        self.gamma = gamma
+        self.vector_dim = vector_dim
+        os.makedirs(self.base_dir, exist_ok=True)
+        
+        # File paths
+        self.episodes_path = os.path.join(self.base_dir, 'episodes.jsonl')
+        self.vector_index_path = os.path.join(self.base_dir, 'state_vectors.index')
+        self.step_metadata_path = os.path.join(self.base_dir, 'step_metadata.pkl')
+        
+        self.eval_llm_model = eval_llm_model
+        
+        self.query_vector_cache = []
+        
+        # Episode counter
+        self.current_episode_number = self._get_episode_count()
+    
+        self.encoder = SentenceTransformer(embedding_model)
+        print(f"✓ Loaded embedding model: {embedding_model} (dim: {vector_dim})")
+        # Initialize FAISS vector database
+        self._init_vector_database()
+    
+    def _init_vector_database(self):
+        # Load or create FAISS index
+        if os.path.exists(self.vector_index_path):
+            self.vector_index = faiss.read_index(self.vector_index_path)
+            print(f"✓ Loaded vector index: {self.vector_index.ntotal} vectors")
+        else:
+            # Use Inner Product for cosine similarity (with normalized vectors)
+            self.vector_index = faiss.IndexFlatIP(self.vector_dim)
+            print(f"✓ Created new vector index")
+        
+        # Load metadata
+        if os.path.exists(self.step_metadata_path):
+            with open(self.step_metadata_path, 'rb') as f:
+                self.step_metadata = pickle.load(f)
+            print(f"✓ Loaded metadata: {len(self.step_metadata)} entries")
+        else:
+            self.step_metadata = []
+    
+    def _save_vector_database(self):
+        """Save FAISS index and metadata to disk."""
+        if self.vector_index is not None:
+            faiss.write_index(self.vector_index, self.vector_index_path)
+        if self.step_metadata:
+            with open(self.step_metadata_path, 'wb') as f:
+                pickle.dump(self.step_metadata, f)
+    
+    def _get_episode_count(self) -> int:
+        """Count existing episodes."""
+        if not os.path.exists(self.episodes_path):
+            return 0
+        with open(self.episodes_path, 'r') as f:
+            return sum(1 for _ in f)
+    
+    def _encode_state(self, state_text: str, history_summary: str = ""):        
+        # Combine current state with history context
+        if history_summary:
+            combined_text = f"History: {history_summary}\n\nCurrent: {state_text}"
+        else:
+            combined_text = f"Current: {state_text}"
+        
+        try:
+            # Encode and normalize for cosine similarity
+            embedding = self.encoder.encode([combined_text])[0]
+            embedding = embedding / (np.linalg.norm(embedding) + 1e-8)
+            return embedding.astype('float32')
+        except Exception as e:
+            print(f"Encoding error: {e}")
+            return None
+
+    # def _encode_state(self, state_text: str, location: str = ""):        
+    #     # Combine current state with history context
+    #     if location:
+    #         combined_text = f"location: {location}\n\nCurrent: {state_text}"
+    #     else:
+    #         combined_text = f"Current: {state_text}"
+        
+    #     try:
+    #         # Encode and normalize for cosine similarity
+    #         embedding = self.encoder.encode([combined_text])[0]
+    #         embedding = embedding / (np.linalg.norm(embedding) + 1e-8)
+    #         return embedding.astype('float32')
+    #     except Exception as e:
+    #         print(f"Encoding error: {e}")
+    #         return None
+    
+    def add_episode(
+        self,
+        game_history: List[Dict[str, Any]],
+        final_score: float,
+        state: str,
+        success: bool = False
+    ):
+        self.current_episode_number += 1
+        
+        # Save episode metadata
+        episode_data = {
+            'episode_number': self.current_episode_number,
+            'final_score': final_score,
+            'success': success,
+            'steps': len(game_history),
+            'history': game_history
+        }
+        
+        with open(self.episodes_path, 'a') as f:
+            f.write(json.dumps(episode_data) + '\n')
+            
+        
+        llm_step_scores, step_reasonings = self.evaluate_step_scores_with_llm(
+                game_history=game_history,
+                state=state,
+                final_score=final_score,
+                success=success,
+                temperature=0.8
+            )
+        try:
+            assert len(llm_step_scores) == len(game_history)
+        except Exception:
+            breakpoint()
+        
+        self._add_steps_to_vector_db(game_history, final_score, llm_step_scores)
+        
+        # Save vector database
+        self._save_vector_database()
+    
+    def _add_steps_to_vector_db(self, game_history: List[Dict], final_score: float, llm_step_scores):
+        n = len(llm_step_scores)
+        
+        for i, step in enumerate(game_history):
+            # Get current state
+            state_text = step.get('state', '')
+            action = step.get('action', '')
+            
+            # Encode state with history
+            vector = self.query_vector_cache[i]
+            if vector is None:
+                continue
+            
+            # Calculate discounted future reward from this step
+            discounted_reward = 0.0
+            for j in range(i, n):
+                if step.get('reward', 0) > 0:
+                    reward = max(llm_step_scores[j], step.get('reward', 0))
+                else:
+                    reward = llm_step_scores[j]
+                discounted_reward += (self.gamma ** (j - i)) * reward
+            
+            # Store metadata for this step
+            metadata = {
+                'episode_number': self.current_episode_number,
+                'step_index': i,
+                'state': state_text,
+                'action': action,
+                'reward': step.get('reward', 0),
+                'score': step.get('score', 0),
+                'llm_step_score': llm_step_scores[i] if i < len(llm_step_scores) else 0,
+                'discounted_reward': discounted_reward,
+                'episode_final_score': final_score
+            }
+            
+            # Add to vector index and metadata list
+            self.vector_index.add(np.array([vector]))
+            self.step_metadata.append(metadata)
+    
+    def retrieve_similar(
+        self,
+        current_state: str,
+        game_history: List[Dict],
+        current_summary: str,
+        k: int = 5,
+        r: float = 0.9
+    ):  
+        # Encode current state with history
+        query_vector = self._encode_state(current_state, current_summary)
+        
+        if query_vector is None:
+            return []
+        
+        self.query_vector_cache.append(query_vector)
+        
+        if self.vector_index is None:
+            print("Dual vector database not available")
+            return []
+        if self.vector_index.ntotal == 0:
+            print("Dual vector database is empty")
+            return []
+        
+        # Phase 1: Recall top candidates using vector similarity (fast)
+        # Use 10x more candidates than final k to ensure good coverage
+        recall_size = min(max(k * 10, 100), self.vector_index.ntotal)
+        
+        scores, indices = self.vector_index.search(query_vector.reshape(1, -1), recall_size)
+        
+        high_similarity_trajectories = []  # similarity > r
+        medium_similarity_trajectories = []  # similarity <= r
+        
+        for score, idx in zip(scores[0], indices[0]):
+            if idx < 0 or idx >= len(self.step_metadata):
+                continue
+            
+            metadata = self.step_metadata[idx]
+            
+            if score > r:
+                high_similarity_trajectories.append((score, metadata["discounted_reward"], metadata))
+            else:
+                medium_similarity_trajectories.append((score, metadata["discounted_reward"], metadata))
+        
+        # Sort both lists by similarity first, then by discounted reward
+        high_similarity_trajectories.sort(key=lambda x: (x[0], x[1]), reverse=True)
+        medium_similarity_trajectories.sort(key=lambda x: (x[0], x[1]), reverse=True)
+        
+        # Combine results: first all high similarity (>0.98), then top k from medium similarity if needed
+        filtered_trajectories = high_similarity_trajectories[:]
+        if len(filtered_trajectories) < k:
+            remaining_needed = k - len(filtered_trajectories)
+            filtered_trajectories.extend(medium_similarity_trajectories[:remaining_needed])
+        
+        return [x[2] for x in filtered_trajectories[:k]]
+    
+    def evaluate_step_scores_with_llm(self, game_history, state, final_score, success, temperature=0.8):
+        """
+        Use LLM to assign scores for each step in the trajectory.
+
+        Args:
+            game_history: List of game history records
+            final_score: Final score
+            success: Whether successful
+            llm_model: LLM model to use
+            temperature: LLM temperature parameter
+
+        Returns:
+            list: Score list for each step
+        """
+        if not game_history:
+            return []
+
+        trajectory_text = ""
+        for i, entry in enumerate(game_history):
+            trajectory_text += f"Step {i}:\n"
+            trajectory_text += f"Current state: {entry.get('old_state', '')}\n"
+            if entry.get('action'):
+                trajectory_text += f"Action: {entry.get('action')}\n"
+                trajectory_text += f"Result state: {entry.get('state', '')}\n"
+            score = entry.get('score', 0)
+            reward = entry.get('reward', 0)
+            trajectory_text += f"Action reward: {reward}\n"
+            trajectory_text += "\n"
+        trajectory_text += f"Step: {len(game_history)}:\nState: {state}\n"
+        sys_prompt = """You are scoring game actions to build training data for future gameplay.
+
+PURPOSE: Rate each action based on its overall impact - positive scores for actions worth repeating, negative for actions to avoid.
+
+SCORING RULES:
+- Positive: Action led to progress, a reward, useful discoveries or new locations discovered
+- Negative: Action wasted time, caused loops, or had no benefit
+- Magnitude: Match the game's typical reward scale (calibrate based on rewards shown in trajectory)
+- Evaluation: Judge by full consequence chain, not just immediate result
+    """
+
+        user_prompt = f"""Score each action in this game session.
+
+Final Result: {"SUCCESS" if success else "FAILURE"}, Final Score: {final_score}
+
+Trajectory:
+{trajectory_text}
+
+JSON FORMAT:
+{{
+"step_analysis": [
+    {{
+    "step": 0,
+    "action": "exact action taken",
+    "detailed_reasoning": "What happened after this action and its consequences? Why should this be repeated or avoided?",
+    "score": 5,
+}}
+],
+}}
+
+Provide complete JSON response:"""
+
+        try:
+            messages = [
+                {"role": "system", "content": sys_prompt},
+                {"role": "user", "content": user_prompt},
+            ]
+            
+            if USE_LOCAL_MODEL and _local_pipeline is not None:
+                outputs = _local_pipeline(
+                    messages,
+                    max_new_tokens=max_tokens,
+                    temperature=temperature,
+                    do_sample=True,
+                )
+                scores_text = outputs[0]["generated_text"][-1]["content"]
+
+            response = LLM_CLIENT.chat.completions.create(
+                model=self.eval_llm_model,
+                messages=messages,
+                temperature=temperature,
+                max_tokens=8192,
+            )
+            scores_text = response.choices[0].message.content.strip()
+        except Exception as e:
+            print(getattr(e, 'response', None) and e.response.text)
+            print(traceback.format_exc())
+            print(f"[Error] LLM call failed: {e}")
+            print("[Warning] Falling back to reward-based scoring")
+            return [entry.get('reward', 0) for entry in game_history], ["No reasoning available" for _ in game_history]
+
+        if scores_text:
+            if scores_text.startswith('```json'):
+                scores_text = scores_text.replace('```json', '').replace('```', '').strip()
+            elif scores_text.startswith('```'):
+                scores_text = scores_text.replace('```', '').strip()
+
+            try:
+                response_data = json.loads(scores_text)
+            except json.JSONDecodeError as e:
+                print(f"[Warning] JSON decode error: {e}")
+                print(f"[Warning] Problematic JSON text (first 500 chars): {scores_text[:500]}")
+                print(f"[Warning] Falling back to reward-based scoring")
+                return [entry.get('reward', 0) for entry in game_history], ["No reasoning available" for _ in game_history]
+
+            if isinstance(response_data, dict):
+                if 'step_analysis' in response_data:
+                    step_analyses = response_data['step_analysis']
+
+                    analysis_dict = {}
+                    for analysis in step_analyses:
+                        step_num = analysis.get('step', -1)
+                        if step_num >= 0:
+                            analysis_dict[step_num] = analysis
+
+                    scores = []
+                    reasonings = []
+                    missing_steps = []
+
+                    print("\n=== LLM Step Analysis ===")
+                    for i in range(len(game_history)):
+                        if i in analysis_dict:
+                            analysis = analysis_dict[i]
+                            step_score = analysis.get('score', 0)
+                            reasoning = analysis.get('detailed_reasoning', 'No reasoning')
+                            scores.append(step_score)
+                            reasonings.append(reasoning)
+                            print(f"Step {i}: {analysis.get('action', 'Unknown')}")
+                            print(f"  Reasoning: {reasoning}")
+                            print(f"  Score: {step_score}")
+                            print()
+                        else:
+                            scores.append(0)
+                            reasonings.append('Missing from LLM analysis - filled with score 0')
+                            missing_steps.append(i)
+                            print(f"Step {i}: [MISSING FROM LLM ANALYSIS - FILLED WITH 0]")
+                            print(f"  Action: {game_history[i].get('action', 'Unknown')}")
+                            print(f"  Score: 0 (auto-filled)")
+                            print()
+
+                    if missing_steps:
+                        print(f"[Warning] LLM missed analyzing steps: {missing_steps}")
+                        print(f"[Info] Filled missing steps with score 0")
+
+                    if 'overall_assessment' in response_data:
+                        print(f"Overall Assessment: {response_data['overall_assessment']}")
+                        print("=" * 50)
+                else:
+                    print(f"[Warning] No step_analysis found in JSON format: {response_data}")
+                    return [entry.get('reward', 0) for entry in game_history], ["No reasoning available" for _ in game_history]
+            elif isinstance(response_data, list):
+                scores = response_data
+                reasonings = ["Legacy format - no reasoning available" for _ in game_history]
+            else:
+                print(f"[Warning] Unexpected JSON format: {response_data}")
+                return [entry.get('reward', 0) for entry in game_history], ["No reasoning available" for _ in game_history]
+
+
+            return scores, reasonings
+        else:
+            print("[Warning] No valid response from LLM for step scoring")
+            return [entry.get('reward', 0) for entry in game_history], ["No reasoning available" for _ in game_history]

mcp_server.py

@@ -26,6 +26,8 @@ Then open the MCP Inspector in your browser to test the tools interactively.
 
 import sys
 import os
+import json
+from dataclasses import asdict
 
 # Add parent directory to path to import games module
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -60,9 +62,9 @@ class GameManager:
         self.state = None
         self.game_name: str = ""
         # TODO: Add more state tracking
-        # self.history: list[tuple[str, str]] = []
-        # self.explored_locations: dict[str, set[str]] = {}
-        # self.current_location: str = ""
+        self.history: list[tuple[str, str]] = []
+        self.explored_locations: dict[str, set[str]] = {}
+        self.current_location: str = ""
     
     def initialize(self, game: str = "zork1"):
         """Initialize or reset the game."""
@@ -71,7 +73,21 @@ class GameManager:
         self.state = self.env.reset()
         # TODO: Reset your state tracking here
         return self.state.observation
-    
+
+    def get_full_state(self, include_valid_actions: bool = True) -> dict:
+        """
+        Convert current env GameState (dataclass from TextAdventureEnv) to a JSON-ready dict.
+        """
+        data = asdict(self.state) 
+
+        if include_valid_actions:
+            try:
+                data["valid_actions"] = self.env.get_valid_actions()  
+            except Exception:
+                data["valid_actions"] = []
+
+        return data
+        
     def step(self, action: str) -> str:
         """Execute an action and return the result."""
         if self.env is None:
@@ -135,12 +151,15 @@ def play_action(action: str) -> str:
     # TODO: You might want to add action validation here
     # TODO: You might want to include score changes in the response
     
-    result = game.step(action)
+    game.step(action)
     
     # Optional: Append score info
     # result += f"\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"
     
-    return result
+    payload = {
+        "state": game.get_full_state(include_valid_actions=True),
+    }
+    return json.dumps(payload)
 
 
 # TODO: Implement additional tools to help your agent
@@ -158,17 +177,17 @@ def play_action(action: str) -> str:
 #     pass
 
 
-# @mcp.tool()
-# def inventory() -> str:
-#     """
-#     Check what the player is carrying.
-#     
-#     Returns:
-#         List of items in the player's inventory
-#     """
-#     game = get_game()
-#     result = game.step("inventory")
-#     return result
+@mcp.tool()
+def inventory() -> str:
+    """
+    Check what the player is carrying.
+    
+    Returns:
+        List of items in the player's inventory
+    """
+    game = get_game()
+    result = game.step("inventory")
+    return result
 
 
 # @mcp.tool()
@@ -184,21 +203,52 @@ def play_action(action: str) -> str:
 #     pass
 
 
-# @mcp.tool()
-# def get_valid_actions() -> str:
-#     """
-#     Get a list of likely valid actions from the current location.
-#     
-#     Returns:
-#         List of actions that might work here
-#     """
-#     # This is a hint: Jericho provides get_valid_actions()
-#     game = get_game()
-#     if game.env and game.env.env:
-#         valid = game.env.env.get_valid_actions()
-#         return "Valid actions: " + ", ".join(valid[:20])
-#     return "Could not determine valid actions"
+@mcp.tool()
+def get_valid_actions() -> str:
+    """
+    Get a list of likely valid actions from the current location.
+    
+    Returns:
+        List of actions that might work here
+    """
+    # This is a hint: Jericho provides get_valid_actions()
+    game = get_game()
+    if game.env and game.env.env:
+        valid = game.env.env.get_valid_actions()
+        return "Valid actions: " + ", ".join(valid[:20])
+    return "Could not determine valid actions"
+
+@mcp.tool()
+def reset_game() -> str:
+    """
+    Reset the game to the beginning or switch to a different game.
+    
+    Use this to start over if you get stuck, die, or want to try a different game.
+    
+    Args:
+        game: Game name (e.g., 'zork1', 'zork2', 'advent', 'enchanter')
+              Use list_games() to see available options.
+    
+    Returns:
+        The initial game text
+    """
+    global _game
+    try:
+        if _game.env is None:
+            game = get_game()
+        else:
+            game_name = os.environ.get("GAME", "zork1")
+            _game.initialize(game_name)
+            game = _game
+        
+        payload = {
+            "state": game.get_full_state(include_valid_actions=True),
+        }
+        return json.dumps(payload)
+    
+    except ValueError as e:
 
+        return json.dumps({"error": str(e)})
 
 # =============================================================================
 # Run the server

requirements.txt

@@ -7,3 +7,9 @@
 # Add any additional packages your agent needs below:
 # numpy
 # requests
+
+python-dotenv
+spacy
+
+faiss
+sentence-transformers