Branch used to run the experiments.

MANIFEST.in

@@ -2,6 +2,7 @@ include requirements.txt
 include README.md
 include LICENSE
 include pyproject.toml
+include tales/jericho/games.json
 
 global-exclude */__pycache__/*
 

agents/llm.py

@@ -102,6 +102,7 @@ def act(self, obs, reward, done, infos):
             "claude-3.5-haiku",
             "claude-3.5-sonnet",
             "claude-3.5-sonnet-latest",
+            "claude-3.7-sonnet",
         ]:
             # For these models, we cannot set the seed.
             llm_kwargs.pop("seed")
@@ -120,9 +121,12 @@ def act(self, obs, reward, done, infos):
         stats = {
             "prompt": format_messages_to_markdown(messages),
             "response": response.text(),
-            "nb_tokens": self.token_counter(messages=messages, text=response.text()),
+            "nb_tokens_prompt": self.token_counter(messages=messages),
+            "nb_tokens_response": self.token_counter(text=response.text()),
         }
 
+        stats["nb_tokens"] = stats["nb_tokens_prompt"] + stats["nb_tokens_response"]
+
         return action, stats
 
     def build_messages(self, observation):

agents/reasoning.py

@@ -40,6 +40,8 @@ def __init__(self, *args, **kwargs):
             "o1-mini",
             "o1-preview",
             "o3-mini",
+            "o4-mini",
+            "o3",
         ]
 
         # Provide the API key, if one is needed and has been provided
@@ -91,9 +93,13 @@ def params(self):
         stop=stop_after_attempt(100),
     )
     def _llm_call_from_conversation(self, conversation, *args, **kwargs):
-        response = conversation.prompt(*args, **kwargs)
-        response.duration_ms()  # Forces the response to be computed.
-        return response
+        for i in range(10):
+            response = conversation.prompt(*args, **kwargs)
+            response.duration_ms()  # Forces the response to be computed.
+            if response.text():
+                return response  # Non-empty response, otherwise retry.
+
+        return ""
 
     def _llm_call_from_messages(self, messages, *args, **kwargs):
         conversation = messages2conversation(self.model, messages)
@@ -116,10 +122,30 @@ def act(self, obs, reward, done, infos):
             else:
                 llm_kwargs["max_tokens"] = self.reasoning_effort
 
-        elif self.llm in ["o1", "o1-preview", "o3-mini"]:
+        elif self.llm in [
+            "o1",
+            "o1-preview",
+            "o3-mini",
+            "o4-mini",
+            "o3",
+            "gpt-5",
+            "gpt-5-mini",
+            "gpt-5-nano",
+        ]:
             llm_kwargs["reasoning_effort"] = self.reasoning_effort
 
-        if self.llm in ["o1", "o1-mini", "o1-preview", "o3-mini", "claude-3.7-sonnet"]:
+        if self.llm in [
+            "o1",
+            "o1-mini",
+            "o1-preview",
+            "o3-mini",
+            "o4-mini",
+            "o3",
+            "claude-3.7-sonnet",
+            "gpt-5",
+            "gpt-5-mini",
+            "gpt-5-nano",
+        ]:
             # For these models, we cannot set the temperature.
             llm_kwargs.pop("temperature")
 
@@ -137,10 +163,43 @@ def act(self, obs, reward, done, infos):
         messages = self.build_messages(f"{obs}\n> ")
         response = self._llm_call_from_messages(messages, **llm_kwargs)
         response_text = response.text()
-
         action = response.text().strip()
 
+        if action == "":
+            # If the action is empty, we need to retry.
+            action = "(empty)"
+
         thinking = None
+        if "Qwen3" in self.llm:
+            # Strip the reasoning <think> and </think>.
+            reasoning_end = action.find("</think>")
+            if reasoning_end == -1:
+                # Send another request to get the action with the current reasoning.
+                messages.append(
+                    {
+                        "role": "assistant",
+                        "content": response_text.strip() + "</think>",
+                    }
+                )
+                llm_kwargs["max_tokens"] = (
+                    100  # Text actions should be short phrases but deepseek forces thought process by starting the generation with <think>.
+                )
+                llm_kwargs["temperature"] = self.act_temp
+                llm_kwargs["extra_body"] = {
+                    "chat_template_kwargs": {"enable_thinking": False}
+                }
+                response = self._llm_call_from_messages(messages, **llm_kwargs)
+                response_text += "</think>" + response.text()
+                action = response_text.strip()
+                reasoning_end = action.find("</think>") + len("</think>")
+            else:
+                reasoning_end += len("</think>")
+
+            # Extract the reasoning part from the response.
+            thinking = action[:reasoning_end].strip()
+            # Extract the action part from the response.
+            action = action[reasoning_end:].strip()
+
         if "DeepSeek-R1" in self.llm:
             # Strip the reasoning <think> and </think>.
             reasoning_end = action.find("</think>")
@@ -192,11 +251,48 @@ def act(self, obs, reward, done, infos):
             "prompt": format_messages_to_markdown(messages),
             "thinking": thinking,
             "response": response_text,
-            "nb_tokens": self.token_counter(messages=messages, text=response_text),
         }
 
-        if thinking is not None:
-            stats["nb_tokens"] += self.token_counter(text=thinking)
+        if self.llm in ["gemini-2.5-pro-preview-03-25", "gemini-2.5-pro-preview-05-06"]:
+            stats["nb_tokens_prompt"] = response.usage().input
+            stats["nb_tokens_thinking"] = response.usage().details.get(
+                "thoughtsTokenCount", 0
+            )
+            stats["nb_tokens_response"] = response.usage().output
+
+        elif self.llm in [
+            "o1",
+            "o1-mini",
+            "o1-preview",
+            "o3-mini",
+            "o4-mini",
+            "o3",
+            "gpt-5",
+            "gpt-5-mini",
+            "gpt-5-nano",
+        ]:
+            # stats["nb_tokens_prompt"] = self.token_counter(messages=messages),
+            # stats["nb_tokens_response"] = self.token_counter(text=response_text)
+            stats["nb_tokens_prompt"] = response.usage().input
+            stats["nb_tokens_response"] = response.usage().output
+            # For these models, we need to look at the API response
+            # stats["nb_tokens_thinking"] = response.usage().details["completion_tokens_details"]["reasoning_tokens"]
+            stats["nb_tokens_thinking"] = response.response_json["usage"][
+                "completion_tokens_details"
+            ]["reasoning_tokens"]
+
+        else:
+            stats["nb_tokens_prompt"] = (self.token_counter(messages=messages),)
+            stats["nb_tokens_thinking"] = (
+                self.token_counter(text=thinking) if thinking else 0
+            )
+            stats["nb_tokens_response"] = self.token_counter(text=response_text)
+
+        stats["nb_tokens"] = (
+            stats["nb_tokens_prompt"]
+            + stats["nb_tokens_response"]
+            + stats["nb_tokens_thinking"]
+        )
 
         return action, stats
 

benchmark.py

@@ -187,14 +187,17 @@ def evaluate(agent, env_name, args):
                     "episode/normalized_score": norm_score,
                     "episode/normalized_highscore": norm_highscore,
                     "episode/token_usage": stats["nb_tokens"],
+                    "episode/token_usage_thinking": stats.get("nb_tokens_thinking", 0),
                 },
                 step=step,
             )
 
             # fmt: off
             results.append([
                 step, score, max_score, norm_score, moves,
-                prev_obs, action, feedback, stats["prompt"], stats["response"], stats.get("thinking"), stats["nb_tokens"]
+                prev_obs, action, feedback,
+                stats["prompt"], stats["response"], stats.get("thinking"),
+                stats["nb_tokens"], stats["nb_tokens_prompt"], stats["nb_tokens_response"], stats.get("nb_tokens_thinking", 0),
             ])
             # fmt: on
 
@@ -256,7 +259,9 @@ def evaluate(agent, env_name, args):
     # fmt: off
     columns = [
         "Step", "Score", "Max Score", "Normalized Score", "Moves",
-        "Observation", "Action", "Feedback", "Prompt", "Response", "Thinking", "Token Usage"
+        "Observation", "Action", "Feedback",
+        "Prompt", "Response", "Thinking",
+        "Token Usage", "Prompt Tokens", "Response Tokens", "Thinking Tokens",
     ]
     # fmt: on
     df = pd.DataFrame(results, columns=columns)
@@ -270,6 +275,9 @@ def evaluate(agent, env_name, args):
         "total/Wins": stats["nb_wins"],
         "total/Resets": stats["nb_resets"],
         "total/Tokens": df["Token Usage"].sum(),
+        "total/Prompt Tokens": df["Prompt Tokens"].sum(),
+        "total/Response Tokens": df["Response Tokens"].sum(),
+        "total/Thinking Tokens": df["Thinking Tokens"].sum(),
         "final/Highscore": stats["highscore"],
         "final/Game Max Score": stats["max_score"],
         "final/Normalized Score": stats["norm_score"],

pyproject.toml

@@ -4,11 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tale-suite"
-version = "1.0.0rc1"
 description = "TALES: Text-Adventure Learning Environment Suite"
 readme = "README.md"
 requires-python = ">=3.12"
-dynamic = ["dependencies"]
+dynamic = ["dependencies", "version"]
 
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -18,6 +17,7 @@ classifiers = [
 
 [tool.setuptools.dynamic]
 dependencies = {file = ["requirements.txt"]}
+version = {attr = "tales.version.__version__"}
 
 [tool.setuptools.packages.find]
 exclude = ["wandb/*", "logs/*", "website/*"]

requirements.txt

@@ -13,13 +13,12 @@ wandb
 numpy
 pandas
 
-# llm>=0.18.0
-llm @ git+https://github.com/MarcCote/llm.git@add_extra_body_option
+llm>=0.27.1
 llm-anthropic
 llm-gemini
-llm-azure-openai @ git+https://github.com/MarcCote/llm-azure-openai.git@generic_ad_auth
+llm-azure-openai
 anthropic
 google-genai
-tiktoken
+tiktoken>=0.11.0
 tenacity
 transformers