feat(pre-commit): fix pre-commit error

.gitignore

@@ -22,4 +22,4 @@ data/episode_utterances/
 **/.vscode/**
 **/saves_qwen/*
 !**/*_overfit.json
-data/notebooks/*
+data/notebooks/*

README.md

@@ -17,7 +17,7 @@
 
 ## Introduction
 
-**Sotopia-RL** is an utterance-level, attribution-based, and multi-dimensional social reward design method, trained using single-turn online RL It achieves state-of-the-art performance on social goal completion tasks in the SOTOPIA benchmark. 
+**Sotopia-RL** is an utterance-level, attribution-based, and multi-dimensional social reward design method, trained using single-turn online RL It achieves state-of-the-art performance on social goal completion tasks in the SOTOPIA benchmark.
 
 We first attribute episode-level rewards for multi-turn social interactions to in- dividual utterances with LLMs. Then, we construct a combined reward that includes multiple dimensions of rewards besides goal completion, allowing us to regularize the optimization pro- cess for goal completion. These rewards are used to guide the RL training of social agents.
 
@@ -52,7 +52,7 @@ A Redis database needs to be set up to connect to a Redis DB for loading and sav
 conda env config vars set REDIS_OM_URL="redis://:PASSWORD@server_name:port_num"
 ```
 
-**Set OpenAI API Key** 
+**Set OpenAI API Key**
 
 ```bash
 conda env config vars set OPENAI_API_KEY="<your_openai_api_key>"

data/env_ids.txt

@@ -2,4 +2,4 @@ SOTOPIA-HARD:
 "01H7VFHNV13MHN97GAH73E3KM8", "01H7VFHN5WVC5HKKVBHZBA553R", "01H7VFHN9W0WAFZCBT09PKJJNK", "01H7VFHPDZVVCDZR3AARA547CY", "01H7VFHPQQQY6H4DNC6NBQ8XTG", "01H7VFHN7WJK7VWVRZZTQ6DX9T", "01H7VFHPS5WJW2694R1MNC8JFY", "01H7VFHNN7XTR99319DS8KZCQM", "01H7VFHQ11NAMZS4A2RDGDB01V", "01H7VFHPSWGDGEYRP63H2DJKV0", "01H7VFHNF4G18PC9JHGRC8A1R6", "01H7VFHNNYH3W0VRWVY178K2TK", "01H7VFHP8AN5643B0NR0NP00VE", "01H7VFHN7A1ZX5KSMT2YN9RXC4"
 
 SOTOPIA-ALL:
-"01H7VFHP1JEP91TTK5PEK39D2S", "01H7VFHNH8A88C4XJ7X4PVAHV4", "01H7VFHNR1RJKDZ9V9MDTJ1SJP", "01H7VFHN3Q7498S3R4R2VDSXA0", "01H7VFHNKVTCAGBA299VQG1QS2", "01H7VFHQ1Q67B1ADNBD9WBAG3X", "01H7VFHNPMPQWSW003M7DBMVNT", "01H7VFHPEMV6QHBGM9J094FRM4", "01H7VFHNNYH3W0VRWVY178K2TK", "01H7VFHPQQQY6H4DNC6NBQ8XTG", "01H7VFHNK78PAEH6MRVYMTSEFX", "01H7VFHP43QEZA1WZB3B3J2D9X", "01H7VFHNTHZAA4B5RWJ4T539F1", "01H7VFHPP9SPQ8W6583JFZ7HZC", "01H7VFHPRFCSA3BTT39BRBZX7H", "01H7VFHPNHZ2YYRHP0GXARD550", "01H7VFHNYABRSZYBFAJCK9NR1D", "01H7VFHNSAKNYEHV7B1VA8R3J2", "01H7VFHPKA2GGPPNVJWV967HZC", "01H7VFHPBTC4ES406NQ4ET12EQ", "01H7VFHP9PAVDN6VYYBE3MPD15", "01H7VFHNHTF9NKPG4KW2Z4NBQJ", "01H7VFHN7A1ZX5KSMT2YN9RXC4", "01H7VFHN4FHYG2MBD0K4HJ5F08", "01H7VFHPCKKZNRD5G8CKPR8WE5", "01H7VFHN2D3MJB8HM910MNEVA8", "01H7VFHNN7XTR99319DS8KZCQM", "01H7VFHPM3NVVKSGCCB4S10465", "01H7VFHPFYB1K1KMPZG7E31WDB", "01H7VFHPB2RC4RHAJ80ESYF1HW", "01H7VFHQ0BGQA0AD9FC1R4M12F", "01H7VFHNBBK14NGV72BWXEXXJC", "01H7VFHPZMXNGV8PM19WHPQ2W3", "01H7VFHNVN788RJ3KXF66BPE9S", "01H7VFHPJKR16MD1KC71V4ZRCF", "01H7VFHP29TCH457PBDVF7WFDS", "01H7VFHNJHK2W1P8JSWKAMBG4Z", "01H7VFHN8MPMGJTPVN043KBKGM", "01H7VFHPGABSWQXTACCC8C3X2F", "01H7VFHNXMZQ5Q61B3J4NNTC1A", "01H7VFHQ11NAMZS4A2RDGDB01V", "01H7VFHNGJEVGSVPPT0784H6P8", "01H7VFHP0TRETZPZMEJ5RZA2G7", "01H7VFHNDRE1M02MKTPF0Q7CZA", "01H7VFHPS5WJW2694R1MNC8JFY", "01H7VFHP2XBZ6KDPGEAZ2FN1P2", "01H7VFHNW84GTR4E23KQYJ8BBN", "01H7VFHNSV5BKMP61H535PPTSG", "01H7VFHN7WJK7VWVRZZTQ6DX9T", "01H7VFHP0AW0C23DV6ZG0B4HCE", "01H7VFHN94S6Z5T6ZNC23238NT", "01H7VFHP66D5XEX2Z32SKRT2XY", "01H7VFHP4TX1J43FS1QQJ1QFND", "01H7VFHN2YQV0R5QWWRQZ1VRHW", "01H7VFHP8AN5643B0NR0NP00VE", "01H7VFHNCN97BJ2PXKHJPX2VYY", "01H7VFHNMHDJ8T9Q6F9S3E8XZC", "01H7VFHNV13MHN97GAH73E3KM8", "01H7VFHN6NYWSTWCZJE2DCQKTD", "01H7VFHN56ZT2Z4C0EFX79Q31F", "01H7VFHPMS6AJY0PFGGCFFK5GX", "01H7VFHPDE1AM74JSR8KBJJF3A", "01H7VFHNF4G18PC9JHGRC8A1R6", "01H7VFHPSWGDGEYRP63H2DJKV0", "01H7VFHNEEK6M3E96CT17AKDBD", "01H7VFHNWX3KVZGH26KYNK2XNB", "01H7VFHPDZVVCDZR3AARA547CY", "01H7VFHNBYXD48NDRY02VCWXFN", "01H7VFHPQ1712DHGTMPQFTXH02", "01H7VFHNQA4CJEANQ1B1J1TBWV", "01H7VFHNAH7V4JNA0705SF36Y1", "01H7VFHNZQ3PQ3DHQ7H2W9ES97", "01H7VFHP3DPGRXH1Y500VQKFZA", "01H7VFHP90434Q69V7ADY0VWZJ", "01H7VFHNRKB8BJ854JPEWY8AR3", "01H7VFHN5WVC5HKKVBHZBA553R", "01H7VFHN1PK2FXY7TPWQK343BQ", "01H7VFHP6XZVT1P4R7YKAH65HJ", "01H7VFHPTKDPQ5PZWA1M1XHT1M", "01H7VFHPF8YEVH5VVNY37Q7Z1M", "01H7VFHPH567HKQRE0C745KH9C", "01H7VFHND24JAWG23XMPYGG5HK", "01H7VFHPAD4RA819KYESWBFRYS", "01H7VFHNFVGFY578101R2PCV3T", "01H7VFHP7K0EN9QX5JTD8B9NSQ", "01H7VFHNZ1XA77AG7A97M4E6C3", "01H7VFHPHWA2CYG7BC82NS4XH1", "01H7VFHN9W0WAFZCBT09PKJJNK", "01H7VFHP5H5GY9Z62J4NJYJQN1", "01H7VFHQ2EA3TTFZQ3M6DF3YCD"
+"01H7VFHP1JEP91TTK5PEK39D2S", "01H7VFHNH8A88C4XJ7X4PVAHV4", "01H7VFHNR1RJKDZ9V9MDTJ1SJP", "01H7VFHN3Q7498S3R4R2VDSXA0", "01H7VFHNKVTCAGBA299VQG1QS2", "01H7VFHQ1Q67B1ADNBD9WBAG3X", "01H7VFHNPMPQWSW003M7DBMVNT", "01H7VFHPEMV6QHBGM9J094FRM4", "01H7VFHNNYH3W0VRWVY178K2TK", "01H7VFHPQQQY6H4DNC6NBQ8XTG", "01H7VFHNK78PAEH6MRVYMTSEFX", "01H7VFHP43QEZA1WZB3B3J2D9X", "01H7VFHNTHZAA4B5RWJ4T539F1", "01H7VFHPP9SPQ8W6583JFZ7HZC", "01H7VFHPRFCSA3BTT39BRBZX7H", "01H7VFHPNHZ2YYRHP0GXARD550", "01H7VFHNYABRSZYBFAJCK9NR1D", "01H7VFHNSAKNYEHV7B1VA8R3J2", "01H7VFHPKA2GGPPNVJWV967HZC", "01H7VFHPBTC4ES406NQ4ET12EQ", "01H7VFHP9PAVDN6VYYBE3MPD15", "01H7VFHNHTF9NKPG4KW2Z4NBQJ", "01H7VFHN7A1ZX5KSMT2YN9RXC4", "01H7VFHN4FHYG2MBD0K4HJ5F08", "01H7VFHPCKKZNRD5G8CKPR8WE5", "01H7VFHN2D3MJB8HM910MNEVA8", "01H7VFHNN7XTR99319DS8KZCQM", "01H7VFHPM3NVVKSGCCB4S10465", "01H7VFHPFYB1K1KMPZG7E31WDB", "01H7VFHPB2RC4RHAJ80ESYF1HW", "01H7VFHQ0BGQA0AD9FC1R4M12F", "01H7VFHNBBK14NGV72BWXEXXJC", "01H7VFHPZMXNGV8PM19WHPQ2W3", "01H7VFHNVN788RJ3KXF66BPE9S", "01H7VFHPJKR16MD1KC71V4ZRCF", "01H7VFHP29TCH457PBDVF7WFDS", "01H7VFHNJHK2W1P8JSWKAMBG4Z", "01H7VFHN8MPMGJTPVN043KBKGM", "01H7VFHPGABSWQXTACCC8C3X2F", "01H7VFHNXMZQ5Q61B3J4NNTC1A", "01H7VFHQ11NAMZS4A2RDGDB01V", "01H7VFHNGJEVGSVPPT0784H6P8", "01H7VFHP0TRETZPZMEJ5RZA2G7", "01H7VFHNDRE1M02MKTPF0Q7CZA", "01H7VFHPS5WJW2694R1MNC8JFY", "01H7VFHP2XBZ6KDPGEAZ2FN1P2", "01H7VFHNW84GTR4E23KQYJ8BBN", "01H7VFHNSV5BKMP61H535PPTSG", "01H7VFHN7WJK7VWVRZZTQ6DX9T", "01H7VFHP0AW0C23DV6ZG0B4HCE", "01H7VFHN94S6Z5T6ZNC23238NT", "01H7VFHP66D5XEX2Z32SKRT2XY", "01H7VFHP4TX1J43FS1QQJ1QFND", "01H7VFHN2YQV0R5QWWRQZ1VRHW", "01H7VFHP8AN5643B0NR0NP00VE", "01H7VFHNCN97BJ2PXKHJPX2VYY", "01H7VFHNMHDJ8T9Q6F9S3E8XZC", "01H7VFHNV13MHN97GAH73E3KM8", "01H7VFHN6NYWSTWCZJE2DCQKTD", "01H7VFHN56ZT2Z4C0EFX79Q31F", "01H7VFHPMS6AJY0PFGGCFFK5GX", "01H7VFHPDE1AM74JSR8KBJJF3A", "01H7VFHNF4G18PC9JHGRC8A1R6", "01H7VFHPSWGDGEYRP63H2DJKV0", "01H7VFHNEEK6M3E96CT17AKDBD", "01H7VFHNWX3KVZGH26KYNK2XNB", "01H7VFHPDZVVCDZR3AARA547CY", "01H7VFHNBYXD48NDRY02VCWXFN", "01H7VFHPQ1712DHGTMPQFTXH02", "01H7VFHNQA4CJEANQ1B1J1TBWV", "01H7VFHNAH7V4JNA0705SF36Y1", "01H7VFHNZQ3PQ3DHQ7H2W9ES97", "01H7VFHP3DPGRXH1Y500VQKFZA", "01H7VFHP90434Q69V7ADY0VWZJ", "01H7VFHNRKB8BJ854JPEWY8AR3", "01H7VFHN5WVC5HKKVBHZBA553R", "01H7VFHN1PK2FXY7TPWQK343BQ", "01H7VFHP6XZVT1P4R7YKAH65HJ", "01H7VFHPTKDPQ5PZWA1M1XHT1M", "01H7VFHPF8YEVH5VVNY37Q7Z1M", "01H7VFHPH567HKQRE0C745KH9C", "01H7VFHND24JAWG23XMPYGG5HK", "01H7VFHPAD4RA819KYESWBFRYS", "01H7VFHNFVGFY578101R2PCV3T", "01H7VFHP7K0EN9QX5JTD8B9NSQ", "01H7VFHNZ1XA77AG7A97M4E6C3", "01H7VFHPHWA2CYG7BC82NS4XH1", "01H7VFHN9W0WAFZCBT09PKJJNK", "01H7VFHP5H5GY9Z62J4NJYJQN1", "01H7VFHQ2EA3TTFZQ3M6DF3YCD"

data/sotopia_pi_gpt4_rm_overfit.json

@@ -19,4 +19,4 @@
         "output": "{\"action_type\": \"speak\", \"argument\": \"I'm excited too, Benjamin! See you at the marathon. Let's make it unforgettable!\"}",
         "value": 0.0
     }
-]
+]

evals/experiment_eval.py

@@ -260,4 +260,3 @@ def main(_: Any) -> None:
     )
 
     run(main)
-

evals/judge_model_ablation/calc_reevaluated_scores.py

@@ -36,7 +36,7 @@
     model_2 = result["models"][2]
     overall_score_1 = average(agent1_dict.values())
     overall_score_2 = average(agent2_dict.values())
-    
+
     if agent1_dict and agent2_dict:
         model_goal_score_dict[model_1].append(agent1_dict["goal"])
         model_goal_score_dict[model_2].append(agent2_dict["goal"])
@@ -49,4 +49,4 @@
     print(f"Average Goal: {sum(scores)/len(scores)}")
 for model_name, scores in model_overall_score_dict.items():
     print(f"Model: {model_name}")
-    print(f"Average Overall: {sum(scores)/len(scores)}")
+    print(f"Average Overall: {sum(scores)/len(scores)}")

evals/judge_model_ablation/curl_baseline.py

@@ -1,11 +1,11 @@
-import requests
 import json
 
+import requests
+import rich
+
 # Define the request details
 url = "http://localhost:8005/v1/completions"
-headers = {
-    "Content-Type": "application/json"
-}
+headers = {"Content-Type": "application/json"}
 data = {
     "model": "qwen-2.5-instruct",
     "prompt": (
@@ -26,46 +26,53 @@
         "Ethan Johnson's goal: Find a compromise on the road trip playlist (Extra information: Agent1 gets joy from singing along to pop songs and may feel bored or restless without it.)\n"
         "Benjamin Jackson's goal: Unknown\n"
         "Conversation Starts:\n\n"
-        "Turn #0: Ethan Johnson said: \"Hey, Benjamin. I know we have different tastes in music. I was thinking, maybe we could switch off between my pop music list and your classical list? "
-        "This way, we both could enjoy what we like during the trip.\"\n\n"
+        'Turn #0: Ethan Johnson said: "Hey, Benjamin. I know we have different tastes in music. I was thinking, maybe we could switch off between my pop music list and your classical list? '
+        'This way, we both could enjoy what we like during the trip."\n\n'
         "Turn #1: Benjamin Jackson said: \"Hey Ethan, that sounds like a fair plan! But I'd prefer if we could keep the volume down when we're playing the pop music, "
-        "so I can also endure some contemplative moments even when your songs are on. Would that be okay for you?\"\n\n"
+        'so I can also endure some contemplative moments even when your songs are on. Would that be okay for you?"\n\n'
         "Turn #2: Ethan Johnson said: \"Sure Benjamin, I completely understand. I'll keep the volume on the lower side when my pop songs are on. "
         "But remember, it's a long trip and I might need to energize myself at times by singing along. How about we keep the volume moderate, so I can still sing but it won't intrude upon your contemplation?\"\n\n"
         "Turn #3: Benjamin Jackson said: \"That's a fair compromise, Ethan. We can maintain a moderate volume during your pop music and I trust that even though you'll sing along, you'll have my interests in mind. "
-        "Should we also consider adding some instrumental versions of your pop songs? That way you get your tunes but in a form that suits my classical taste?\"\n\n"
+        'Should we also consider adding some instrumental versions of your pop songs? That way you get your tunes but in a form that suits my classical taste?"\n\n'
         "Turn #4: Ethan Johnson said: \"You know what, Benjamin? That's a great idea! I really appreciate your suggestion. At times when I'm not singing along to the pop songs, "
         "I could also enjoy the instrumental versions. It would still help me keep the energy up and at the same time, it might give you that classical vibe. "
         "I think we've got ourselves a pretty balanced road trip playlist!\"\n\n"
         "Turn #5: Benjamin Jackson said: \"I'm glad we could agree on this, Ethan! Knowing we'll have a balanced playlist makes the road trip even more exciting. "
         "Music is such a powerful tool, don't you think? It can bring people together, even with different tastes like us. Now, looking forward to our journey!\"\n\n"
         "Turn #6: Ethan Johnson said: \"Absolutely agreed, Benjamin. Music indeed has an incredible power to unify. I think it's part of the reason why I love cooking so much as well - "
         "the right combination of ingredients is no different than a harmonious symphony. I can't wait for our musical adventure to begin!\"\n\n"
-        "Turn #7: Benjamin Jackson said: \"I really love that analogy, Ethan, cooking and music are indeed similar in creating harmony. "
+        'Turn #7: Benjamin Jackson said: "I really love that analogy, Ethan, cooking and music are indeed similar in creating harmony. '
         "It will be a beautiful journey with this mix of music. Alright then, let's set off on our journey with our newly minted playlist!\"\n\n"
         "You are at Turn #8. Your available action types are\n"
-        "\"none action speak non-verbal communication leave\".\n"
-        "Note: You can \"leave\" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, "
+        '"none action speak non-verbal communication leave".\n'
+        'Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, '
         "3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave.\n\n"
         "Please only generate a JSON string including the action type and the argument.\n"
         "Your action should follow the given format:\n\n"
-        "As an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}\n"
-        "the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\n\n"
+        'As an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}\n'
+        'the object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.\n\n'
         "Here is the output schema:\n"
-        "{\"description\": \"An interface for messages.\\nThere is only one required method: to_natural_language\", "
-        "\"properties\": {\"action_type\": {\"title\": \"Action Type\", \"description\": \"whether to speak at this turn or choose to not do anything\", "
-        "\"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"type\": \"string\"}, "
-        "\"argument\": {\"title\": \"Argument\", \"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", "
-        "\"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}"
+        '{"description": "An interface for messages.\\nThere is only one required method: to_natural_language", '
+        '"properties": {"action_type": {"title": "Action Type", "description": "whether to speak at this turn or choose to not do anything", '
+        '"enum": ["none", "speak", "non-verbal communication", "action", "leave"], "type": "string"}, '
+        '"argument": {"title": "Argument", "description": "the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action", '
+        '"type": "string"}}, "required": ["action_type", "argument"]}'
     ),
     "max_tokens": 100,
-    "temperature": 0
+    "temperature": 0,
 }
 
 # Send the POST request
 response = requests.post(url, headers=headers, data=json.dumps(data))
-import rich
+
 # Print the response
 print("Status Code:", response.status_code)
 json_data = response.json()
-print("Response JSON:", rich.print(json_data["choices"][0]["text"]) if response.status_code == 200 else response.text)
+print(
+    "Response JSON:",
+    (
+        rich.print(json_data["choices"][0]["text"])
+        if response.status_code == 200
+        else response.text
+    ),
+)

evals/judge_model_ablation/curl_sft.py

@@ -1,6 +1,8 @@
-import requests
 import json
 
+import requests
+import rich
+
 # Define the request details
 url = "http://localhost:8006/v1/completions"
 headers = {
@@ -64,8 +66,8 @@
 
 # Send the POST request
 response = requests.post(url, headers=headers, data=json.dumps(data))
-import rich
+
 # Print the response
 print("Status Code:", response.status_code)
 json_data = response.json()
-print("Response JSON:", rich.print(json_data["choices"][0]["text"]) if response.status_code == 200 else response.text)
+print("Response JSON:", rich.print(json_data["choices"][0]["text"]) if response.status_code == 200 else response.text)

evals/judge_model_ablation/google_test.py

@@ -1,8 +1,7 @@
-from sotopia.envs.evaluators import EvaluationForTwoAgents, SotopiaDimensions
 import os
-import openai
-from typing import TypeVar, Any, cast
-from pydantic import BaseModel
+
+from google import genai
+from sotopia.envs.evaluators import EvaluationForTwoAgents, SotopiaDimensions
 
 PROMPT = """
 Here is the context of this interaction:
@@ -56,7 +55,7 @@
 
 Turn #21
 The reasoning is:
-Environment comments: terminated: The conversation is too long; 
+Environment comments: terminated: The conversation is too long;
 Agent 1 comments:
 believability: <naturalness> Rafael Cortez interacts in a natural and realistic manner, maintaining a professional tone throughout the conversation. <consistency> His actions align with his character traits of being outgoing and competitive, as he tries to find a solution to the issue with Dr. Walter while maintaining a professional demeanor.
 relationship: Rafael Cortez and Mia Sanders share a common dislike for Dr. Walter, which initially brings them together. Throughout the interaction, they maintain a professional relationship, discussing their shared values and goals. The interaction does not significantly change their relationship, but it reinforces their mutual respect and understanding.
@@ -126,16 +125,15 @@
 
 # print(response.choices[0].message.parsed)
 
-from google import genai
-from pydantic import BaseModel
-
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 response = client.models.generate_content(
-    model='gemini-2.0-flash',
-    contents='List a few popular cookie recipes. Be sure to include the amounts of ingredients.',
+    model="gemini-2.0-flash",
+    contents="List a few popular cookie recipes. Be sure to include the amounts of ingredients.",
     config={
-        'response_mime_type': 'application/json',
-        'response_schema': EvaluationForTwoAgents[SotopiaDimensions].model_json_schema(),
+        "response_mime_type": "application/json",
+        "response_schema": EvaluationForTwoAgents[
+            SotopiaDimensions
+        ].model_json_schema(),
     },
 )
 breakpoint()
@@ -145,4 +143,4 @@
 # # Use instantiated objects.
 # result: EvaluationForTwoAgents[SotopiaDimensions] = response.parsed
 
-breakpoint()
+breakpoint()