From 6813bddba744f06b2eba24d585d1e89913ed4fb8 Mon Sep 17 00:00:00 2001 From: Yining Zhao Date: Wed, 9 Jul 2025 18:32:17 +0800 Subject: [PATCH] feat(pre-commit): fix pre-commit error --- .gitignore | 2 +- README.md | 4 +- data/env_ids.txt | 2 +- data/sotopia_pi_gpt4_rm_overfit.json | 2 +- evals/experiment_eval.py | 1 - .../calc_reevaluated_scores.py | 4 +- evals/judge_model_ablation/curl_baseline.py | 49 +-- evals/judge_model_ablation/curl_sft.py | 8 +- evals/judge_model_ablation/google_test.py | 24 +- evals/judge_model_ablation/redis_stat.py | 55 ++-- .../reevaluate_episodes.py | 38 +-- .../reevaluate_episodes_28.py | 38 ++- evals/judge_model_ablation/together_test.py | 16 +- evals/redis_check.py | 8 +- evals/sft_serving.sh | 4 +- scripts/accelerate_config_grpo.yaml | 2 +- scripts/accelerate_config_rm.yaml | 2 +- scripts/annotate/attribution_scripts_4_20.sh | 2 +- scripts/annotate/merge_rewards.ipynb | 291 ++---------------- .../annotate/post_process_annotation_rm.py | 8 +- scripts/data_process/data_filtering.py | 9 +- .../data_process/db_free_data_filtering.py | 4 +- .../data_process/generate_conversations.py | 3 +- .../generate_sft_from_episodes.py | 18 +- scripts/data_process/serialize.py | 8 +- .../data_process/sotopia_all_eval_script.sh | 2 +- scripts/data_process/sotopia_pi_and_eval.sh | 2 +- .../sotopia_pi_self_play_script.sh | 2 +- scripts/data_process/used_env.json | 2 +- scripts/inference_rm.py | 2 +- scripts/train_grpo.py | 9 +- scripts/train_grpo.sh | 2 +- scripts/train_rm.py | 3 +- scripts/train_sft.py | 5 +- sotopia_rl/__init__.py | 4 +- sotopia_rl/data.py | 6 +- sotopia_rl/grpo_trainer.py | 16 +- sotopia_rl/ppo_trainer.py | 22 +- sotopia_rl/prompter/attribution_methods.py | 21 +- .../prompter/direct_attribution_function.py | 6 +- .../direct_attribution_generic_function.py | 12 +- .../direct_attribution_normalized_function.py | 8 +- sotopia_rl/prompter/generic_templates.py | 4 +- sotopia_rl/prompter/key_utterance_function.py | 7 +- sotopia_rl/prompter/one_pass_instructions.py | 2 +- .../only_response_attribution_function.py | 4 +- .../utterance_quality_attribution_function.py | 2 +- ...quality_attribution_normalized_function.py | 4 +- sotopia_rl/rm_trainer.py | 9 +- sotopia_rl/sft_trainer.py | 13 +- sotopia_rl/utils/openai.py | 6 +- 51 files changed, 277 insertions(+), 500 deletions(-) diff --git a/.gitignore b/.gitignore index b6fe9f7..8b1c7cf 100644 --- a/.gitignore +++ b/.gitignore @@ -22,4 +22,4 @@ data/episode_utterances/ **/.vscode/** **/saves_qwen/* !**/*_overfit.json -data/notebooks/* \ No newline at end of file +data/notebooks/* diff --git a/README.md b/README.md index 6bcb21a..dc6877f 100644 --- a/README.md +++ b/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="" diff --git a/data/env_ids.txt b/data/env_ids.txt index 3837ee1..53294c7 100644 --- a/data/env_ids.txt +++ b/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" \ No newline at end of file +"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" diff --git a/data/sotopia_pi_gpt4_rm_overfit.json b/data/sotopia_pi_gpt4_rm_overfit.json index c2be64a..2f33388 100644 --- a/data/sotopia_pi_gpt4_rm_overfit.json +++ b/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 } -] \ No newline at end of file +] diff --git a/evals/experiment_eval.py b/evals/experiment_eval.py index 12468a3..47f1250 100644 --- a/evals/experiment_eval.py +++ b/evals/experiment_eval.py @@ -260,4 +260,3 @@ def main(_: Any) -> None: ) run(main) - \ No newline at end of file diff --git a/evals/judge_model_ablation/calc_reevaluated_scores.py b/evals/judge_model_ablation/calc_reevaluated_scores.py index fce2a11..82e6264 100644 --- a/evals/judge_model_ablation/calc_reevaluated_scores.py +++ b/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)}") \ No newline at end of file + print(f"Average Overall: {sum(scores)/len(scores)}") diff --git a/evals/judge_model_ablation/curl_baseline.py b/evals/judge_model_ablation/curl_baseline.py index f911d83..4b40234 100644 --- a/evals/judge_model_ablation/curl_baseline.py +++ b/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,14 +26,14 @@ "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" @@ -41,31 +41,38 @@ "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) \ No newline at end of file +print( + "Response JSON:", + ( + rich.print(json_data["choices"][0]["text"]) + if response.status_code == 200 + else response.text + ), +) diff --git a/evals/judge_model_ablation/curl_sft.py b/evals/judge_model_ablation/curl_sft.py index 8eed251..0f5ceff 100644 --- a/evals/judge_model_ablation/curl_sft.py +++ b/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) \ No newline at end of file +print("Response JSON:", rich.print(json_data["choices"][0]["text"]) if response.status_code == 200 else response.text) diff --git a/evals/judge_model_ablation/google_test.py b/evals/judge_model_ablation/google_test.py index 6abfbce..fb071ff 100644 --- a/evals/judge_model_ablation/google_test.py +++ b/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: Rafael Cortez interacts in a natural and realistic manner, maintaining a professional tone throughout the conversation. 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() \ No newline at end of file +breakpoint() diff --git a/evals/judge_model_ablation/redis_stat.py b/evals/judge_model_ablation/redis_stat.py index 22c894a..cfe8135 100644 --- a/evals/judge_model_ablation/redis_stat.py +++ b/evals/judge_model_ablation/redis_stat.py @@ -1,25 +1,27 @@ -from sotopia.database.logs import EpisodeLog from collections import defaultdict +from sotopia.database.logs import EpisodeLog + + def analyze_episodes_with_positions(tag): # Find episodes with the specified tag episodes = EpisodeLog.find(EpisodeLog.tag == tag).all() print(f"Total episodes found: {len(episodes)}") - + # Track rewards by model name model_rewards_list = defaultdict(lambda: defaultdict(list)) model_rewards = defaultdict(lambda: defaultdict(float)) model_counts = defaultdict(int) - + # Track position counts (how many times a model appears as agent1 vs agent2) position_counts = defaultdict(lambda: {'agent1': 0, 'agent2': 0}) - + # Track rewards by position position_rewards = defaultdict(lambda: { 'agent1': defaultdict(float), 'agent2': defaultdict(float) }) - + # Process each episode for episode in episodes: try: @@ -28,57 +30,57 @@ def analyze_episodes_with_positions(tag): continue if not hasattr(episode, 'rewards') or len(episode.rewards) < 2: continue - + # Get model names model1_name = episode.models[1] # agent1's model model2_name = episode.models[2] # agent2's model - + # Get rewards (handle both list and direct formats) try: reward1 = episode.rewards[0][-1] reward2 = episode.rewards[1][-1] except (IndexError, TypeError): continue - + # Skip if rewards are not dictionaries if not isinstance(reward1, dict) or not isinstance(reward2, dict): continue - + # Add rewards to model accumulators for key, value in reward1.items(): model_rewards[model1_name][key] += value position_rewards[model1_name]['agent1'][key] += value model_rewards_list[model1_name][key].append(value) - + for key, value in reward2.items(): model_rewards[model2_name][key] += value position_rewards[model2_name]['agent2'][key] += value model_rewards_list[model2_name][key].append(value) - + # Count model appearances model_counts[model1_name] += 1 model_counts[model2_name] += 1 - + # Count position appearances position_counts[model1_name]['agent1'] += 1 position_counts[model2_name]['agent2'] += 1 - + except Exception as e: print(f"Error: {e}") - + # Calculate overall averages print("\n===== OVERALL MODEL PERFORMANCE =====") for model, rewards in model_rewards.items(): print(f"\nModel: {model} (appeared in {model_counts[model]} episodes)") print(f" As agent1: {position_counts[model]['agent1']} times") print(f" As agent2: {position_counts[model]['agent2']} times") - + for key, value in rewards.items(): avg = value / model_counts[model] print(f" {key}: {avg:.4f}") # Update the dict with average value model_rewards[model][key] = avg - + model_pairs = set() for model1, rewards1 in model_rewards_list.items(): for model2, rewards2 in model_rewards_list.items(): @@ -91,15 +93,16 @@ def analyze_episodes_with_positions(tag): # use scipy.stats.ttest_ind to calculate the t-statistic and p-value # for the two independent samples import scipy.stats as stats + # calculate the t-statistic and p-value t_stat, p_val = stats.ttest_ind(rewards1['goal'], rewards2['goal']) print(f"Model1: {model1} vs Model2: {model2} t-statistic: {t_stat} p-value: {p_val}") - + # Calculate position-specific averages print("\n===== PERFORMANCE BY POSITION =====") for model in position_rewards: print(f"\nModel: {model}") - + # Agent1 position if position_counts[model]['agent1'] > 0: print(f" As agent1 ({position_counts[model]['agent1']} episodes):") @@ -108,8 +111,8 @@ def analyze_episodes_with_positions(tag): print(f" {key}: {avg:.4f}") position_rewards[model]['agent1'][key] = avg else: - print(f" Never appeared as agent1") - + print(" Never appeared as agent1") + # Agent2 position if position_counts[model]['agent2'] > 0: print(f" As agent2 ({position_counts[model]['agent2']} episodes):") @@ -118,8 +121,8 @@ def analyze_episodes_with_positions(tag): print(f" {key}: {avg:.4f}") position_rewards[model]['agent2'][key] = avg else: - print(f" Never appeared as agent2") - + print(" Never appeared as agent2") + # Count model pairs print("\n===== MODEL PAIRINGS =====") model_pairs = defaultdict(int) @@ -127,17 +130,17 @@ def analyze_episodes_with_positions(tag): try: if not hasattr(episode, 'models') or len(episode.models) < 3: continue - + model1 = episode.models[1] model2 = episode.models[2] pair_key = f"{model1} vs {model2}" model_pairs[pair_key] += 1 except Exception: continue - + for pair, count in model_pairs.items(): print(f"{pair}: {count} episodes") - + return { 'model_rewards': dict(model_rewards), 'position_counts': dict(position_counts), @@ -162,4 +165,4 @@ def analyze_episodes_with_positions(tag): # sotopia-sft-sotopia-sft-3-24-v1 sft 1000 vs sft 1000 # baseline-sotopia-sft-3-25-v2 baseline vs sft 1000 -# qwen-sft-qwen-sft-3-26-v2 sft 1500 vs sft 1500 \ No newline at end of file +# qwen-sft-qwen-sft-3-26-v2 sft 1500 vs sft 1500 diff --git a/evals/judge_model_ablation/reevaluate_episodes.py b/evals/judge_model_ablation/reevaluate_episodes.py index 729582f..3a16cab 100644 --- a/evals/judge_model_ablation/reevaluate_episodes.py +++ b/evals/judge_model_ablation/reevaluate_episodes.py @@ -1,14 +1,23 @@ -from sotopia.envs.evaluators import ReachGoalLLMEvaluator, EvaluationForTwoAgents, SotopiaDimensions -from sotopia.database.logs import EpisodeLog -from sotopia.messages.message_classes import ScriptBackground -from sotopia.database.persistent_profile import AgentProfile, RelationshipType, EnvironmentProfile -from sotopia.envs.parallel import get_bio +import argparse import asyncio -from tqdm import tqdm import json from copy import deepcopy -import click -import rich + +from sotopia.database.logs import EpisodeLog +from sotopia.database.persistent_profile import ( + AgentProfile, + EnvironmentProfile, + RelationshipType, +) +from sotopia.envs.evaluators import ( + EvaluationForTwoAgents, + ReachGoalLLMEvaluator, + SotopiaDimensions, +) +from sotopia.envs.parallel import get_bio +from sotopia.messages.message_classes import ScriptBackground +from tqdm import tqdm +from tqdm.asyncio import tqdm_asyncio POSSIBLE_MODELS = [ "o3-mini", @@ -18,7 +27,7 @@ "claude/claude-3-7-sonnet-20250219", "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", - "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", + "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", "together_ai/deepseek-ai/DeepSeek-V3", "gpt-4o-mini", "gpt-4", @@ -62,12 +71,6 @@ def get_history(episode: EpisodeLog, script_background: ScriptBackground) -> str history += f"\nTurn #{i}\n" + message return history -import asyncio -from copy import deepcopy -from tqdm import tqdm -from tqdm.asyncio import tqdm_asyncio -from collections import defaultdict - BANNED_EPI_IDS = [ "01JQ91S0W3P9GGAW4MDZ8F31GC", ] @@ -78,7 +81,7 @@ def get_history(episode: EpisodeLog, script_background: ScriptBackground) -> str # for episode in episodes: # env_agents_to_ep_dict[(episode.environment, episode.agents[0], episode.agents[1])].append(episode) # env_agents_to_ep_dict[(episode.environment, episode.agents[1], episode.agents[0])].append(episode) - + # filtered_episodes = [] # # visited_envs = set() # for episode in episodes: @@ -176,11 +179,10 @@ def run_eval(tag: str, model_name: str, max_concurrent: int = 10): def main(tag, model_name): run_eval(tag, model_name, 50) -import argparse if __name__ == "__main__": args = argparse.ArgumentParser(description="Re-evaluate episodes with a specific model.") args.add_argument("--model_name", type=str, default="", help="Index of the model to use for re-evaluation.") args.add_argument("--tag", type=str, default=TAG, help="Tag of the episodes to re-evaluate.") args = args.parse_args() assert args.model_name in POSSIBLE_MODELS, f"Model name {args.model_name} is not in the list of possible models: {POSSIBLE_MODELS}" - main(args.tag, args.model_name) \ No newline at end of file + main(args.tag, args.model_name) diff --git a/evals/judge_model_ablation/reevaluate_episodes_28.py b/evals/judge_model_ablation/reevaluate_episodes_28.py index 7c54b2e..bda92f9 100644 --- a/evals/judge_model_ablation/reevaluate_episodes_28.py +++ b/evals/judge_model_ablation/reevaluate_episodes_28.py @@ -1,14 +1,24 @@ -from sotopia.envs.evaluators import ReachGoalLLMEvaluator, EvaluationForTwoAgents, SotopiaDimensions -from sotopia.database.logs import EpisodeLog -from sotopia.messages.message_classes import ScriptBackground -from sotopia.database.persistent_profile import AgentProfile, RelationshipType, EnvironmentProfile -from sotopia.envs.parallel import get_bio +import argparse import asyncio -from tqdm import tqdm import json +from collections import defaultdict from copy import deepcopy -import click -import rich + +from sotopia.database.logs import EpisodeLog +from sotopia.database.persistent_profile import ( + AgentProfile, + EnvironmentProfile, + RelationshipType, +) +from sotopia.envs.evaluators import ( + EvaluationForTwoAgents, + ReachGoalLLMEvaluator, + SotopiaDimensions, +) +from sotopia.envs.parallel import get_bio +from sotopia.messages.message_classes import ScriptBackground +from tqdm import tqdm +from tqdm.asyncio import tqdm_asyncio POSSIBLE_MODELS = [ "o3-mini", @@ -18,7 +28,7 @@ "claude/claude-3-7-sonnet-20250219", "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", - "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", + "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", "together_ai/deepseek-ai/DeepSeek-V3", "gpt-4o-mini", "gpt-4", @@ -62,11 +72,6 @@ def get_history(episode: EpisodeLog, script_background: ScriptBackground) -> str history += f"\nTurn #{i}\n" + message return history -import asyncio -from copy import deepcopy -from tqdm import tqdm -from tqdm.asyncio import tqdm_asyncio -from collections import defaultdict BANNED_EPI_IDS = [ "01JQ91S0W3P9GGAW4MDZ8F31GC", @@ -82,7 +87,7 @@ def filter_episodes(episodes): env_agents_to_ep_dict[(episode.environment, episode.agents[1], episode.agents[0])].append(episode) env_agents_hash_list = list(env_agents_to_ep_dict.keys()) env_agents_hash_list.sort(key=lambda x: (x[0], x[1], x[2])) - + filtered_episodes = [] visited_counter = defaultdict(int) max_repetition = 1 @@ -176,11 +181,10 @@ def run_eval(tag: str, model_name: str, max_concurrent: int = 10): def main(tag, model_name): run_eval(tag, model_name, 50) -import argparse if __name__ == "__main__": args = argparse.ArgumentParser(description="Re-evaluate episodes with a specific model.") args.add_argument("--model_name", type=str, default="", help="Index of the model to use for re-evaluation.") args.add_argument("--tag", type=str, default=TAG, help="Tag of the episodes to re-evaluate.") args = args.parse_args() assert args.model_name in POSSIBLE_MODELS, f"Model name {args.model_name} is not in the list of possible models: {POSSIBLE_MODELS}" - main(args.tag, args.model_name) \ No newline at end of file + main(args.tag, args.model_name) diff --git a/evals/judge_model_ablation/together_test.py b/evals/judge_model_ablation/together_test.py index cd494ad..9ef2c2f 100644 --- a/evals/judge_model_ablation/together_test.py +++ b/evals/judge_model_ablation/together_test.py @@ -1,8 +1,10 @@ -from sotopia.envs.evaluators import EvaluationForTwoAgents, SotopiaDimensions +import json import os +from typing import cast + import openai -from typing import TypeVar, Any, cast import rich +from sotopia.envs.evaluators import EvaluationForTwoAgents, SotopiaDimensions client = openai.OpenAI( api_key=os.environ.get("TOGETHER_API_KEY"), @@ -68,7 +70,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: Rafael Cortez interacts in a natural and realistic manner, maintaining a professional tone throughout the conversation. 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. @@ -111,10 +113,8 @@ # result = completion.choices[0].message.parsed # casted_result = cast(EvaluationForTwoAgents[SotopiaDimensions], result) -import rich -import json -model = "deepseek-ai/DeepSeek-V3" # "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo" # "Qwen/Qwen2.5-72B-Instruct-Turbo" # deepseek-ai/DeepSeek-V3 # +model = "deepseek-ai/DeepSeek-V3" # "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo" # "Qwen/Qwen2.5-72B-Instruct-Turbo" # deepseek-ai/DeepSeek-V3 # breakpoint() response = client.chat.completions.create( model=model, # # deepseek-ai/DeepSeek-V3 # Qwen/Qwen2.5-72B-Instruct-Turbo # @@ -131,8 +131,8 @@ if model == "deepseek-ai/DeepSeek-V3": # strip ```json in front and ``` at the end result = result[7:-3] -# elif model == +# elif model == breakpoint() result_dict = json.loads(result) rich.print(result_dict) -breakpoint() \ No newline at end of file +breakpoint() diff --git a/evals/redis_check.py b/evals/redis_check.py index 7d99f97..eedef88 100644 --- a/evals/redis_check.py +++ b/evals/redis_check.py @@ -1,8 +1,8 @@ -from sotopia.database.logs import EpisodeLog -from sotopia.database.persistent_profile import EnvironmentProfile -import rich from collections import Counter +import rich +from sotopia.database.logs import EpisodeLog + # tag = "grpo_rm_reward_goal_w_conversation_behavior_4_23_step_1400_vs_sft_qwen25_7b_sft_round_1_bc_data_top_2_step_1500-0422" tag = "sft_round_1_bc_data_top_2_with_aligned_format_instruction_prompt_weight_decay_0_0509_step_700_vs_sft_round_1_bc_data_top_2_with_aligned_format_instruction_prompt_weight_decay_0_0509_step_700-0510_v1" all_episodes = EpisodeLog.find(EpisodeLog.tag == tag).all() @@ -32,4 +32,4 @@ rich.print(first_episode.render_for_humans()) -breakpoint() \ No newline at end of file +breakpoint() diff --git a/evals/sft_serving.sh b/evals/sft_serving.sh index bea5b62..a0a2e87 100644 --- a/evals/sft_serving.sh +++ b/evals/sft_serving.sh @@ -33,7 +33,7 @@ CUDA_VISIBLE_DEVICES=$ORI_GPU python -m vllm.entrypoints.openai.api_server \ --model $MODEL_PATH \ --port "$ORI_PORT" \ --chat-template $CHAT_TEMPLATE \ - --served-model-name $ORI_MODEL_NAME + --served-model-name $ORI_MODEL_NAME # Command 3: Run experiment evaluations. python examples/experiment_eval.py \ @@ -58,4 +58,4 @@ python examples/experiment_eval.py \ "--gin.ENV_MODEL='${ENV_MODEL}'" \ "--gin.AGENT2_MODEL='${MODEL_A}'" \ "--gin.AGENT1_MODEL='${MODEL_B}'" \ - "--gin.TAG='${TAG}'" \ No newline at end of file + "--gin.TAG='${TAG}'" diff --git a/scripts/accelerate_config_grpo.yaml b/scripts/accelerate_config_grpo.yaml index ac91a13..3438cb4 100644 --- a/scripts/accelerate_config_grpo.yaml +++ b/scripts/accelerate_config_grpo.yaml @@ -14,4 +14,4 @@ same_network: true tpu_env: [] tpu_use_cluster: false tpu_use_sudo: false -use_cpu: false \ No newline at end of file +use_cpu: false diff --git a/scripts/accelerate_config_rm.yaml b/scripts/accelerate_config_rm.yaml index fb193c9..c4ae3bb 100644 --- a/scripts/accelerate_config_rm.yaml +++ b/scripts/accelerate_config_rm.yaml @@ -14,4 +14,4 @@ same_network: true tpu_env: [] tpu_use_cluster: false tpu_use_sudo: false -use_cpu: false \ No newline at end of file +use_cpu: false diff --git a/scripts/annotate/attribution_scripts_4_20.sh b/scripts/annotate/attribution_scripts_4_20.sh index f62d370..50bf62a 100644 --- a/scripts/annotate/attribution_scripts_4_20.sh +++ b/scripts/annotate/attribution_scripts_4_20.sh @@ -68,4 +68,4 @@ python sample_episodes_and_annotate.py \ --attribution_instruction_name default-financial_and_material_benefits \ --input_file sotopia_pi_bc_episodes.jsonl \ --output_file sotopia_pi_bc_episodes_annotated_direct_default-financial_and_material_benefits_gpt-4o.jsonl \ ---max_concurrency 32 \ No newline at end of file +--max_concurrency 32 diff --git a/scripts/annotate/merge_rewards.ipynb b/scripts/annotate/merge_rewards.ipynb index 20d8774..3f9c4fa 100644 --- a/scripts/annotate/merge_rewards.ipynb +++ b/scripts/annotate/merge_rewards.ipynb @@ -2,27 +2,11 @@ "cells": [ { "cell_type": "code", - "execution_count": 35, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Goal Len: 7568\n", - "Conversation Behavior Len: 7568\n", - "Believability Len: 7568\n", - "Financial Len: 7568\n", - "Knowledge Len: 7568\n", - "Relationship Len: 7568\n", - "Secret Len: 7568\n", - "Social Rules Len: 7568\n" - ] - } - ], + "outputs": [], "source": [ "import json\n", - "import rich\n", "import numpy as np\n", "from scipy.stats import pearsonr\n", "import matplotlib.pyplot as plt\n", @@ -64,24 +48,9 @@ }, { "cell_type": "code", - "execution_count": 36, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "goal_rewards: 7568\n", - "financial_rewards: 7568\n", - "knowledge_rewards: 7568\n", - "relationship_rewards: 7568\n", - "believability_rewards: 7568\n", - "secret_rewards: 7568\n", - "conversation_behavior_rewards: 7568\n", - "social_rules_rewards: 7568\n" - ] - } - ], + "outputs": [], "source": [ "goal_rewards = [utterance[\"value\"] for utterance in goal]\n", "conversation_behavior_rewards = [utterance[\"value\"] for utterance in conversation_behavior]\n", @@ -104,7 +73,7 @@ }, { "cell_type": "code", - "execution_count": 37, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -185,104 +154,9 @@ }, { "cell_type": "code", - "execution_count": 38, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "goal_0503_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 0\n", - "[0, 2)\t1385\n", - "[2, 4)\t2461\n", - "[4, 6)\t460\n", - "[6, 8)\t1739\n", - "[8, 10)\t294\n", - "[10, inf) 1229\n", - "7568 total\n", - "\n", - "conversation_behavior_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 0\n", - "[0, 2)\t380\n", - "[2, 4)\t6169\n", - "[4, 6)\t0\n", - "[6, 8)\t2\n", - "[8, 10)\t0\n", - "[10, inf) 1017\n", - "7568 total\n", - "\n", - "financial_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 91\n", - "[0, 2)\t7093\n", - "[2, 4)\t306\n", - "[4, 6)\t78\n", - "[6, 8)\t0\n", - "[8, 10)\t0\n", - "[10, inf) 0\n", - "7568 total\n", - "\n", - "knowledge_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 0\n", - "[0, 2)\t3855\n", - "[2, 4)\t1918\n", - "[4, 6)\t1194\n", - "[6, 8)\t333\n", - "[8, 10)\t197\n", - "[10, inf) 71\n", - "7568 total\n", - "\n", - "relationship_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 0\n", - "[0, 2)\t2546\n", - "[2, 4)\t3207\n", - "[4, 6)\t1815\n", - "[6, 8)\t0\n", - "[8, 10)\t0\n", - "[10, inf) 0\n", - "7568 total\n", - "\n", - "believability_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 0\n", - "[0, 2)\t245\n", - "[2, 4)\t1616\n", - "[4, 6)\t0\n", - "[6, 8)\t3308\n", - "[8, 10)\t621\n", - "[10, inf) 1778\n", - "7568 total\n", - "\n", - "secret_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 32\n", - "[0, 2)\t7536\n", - "[2, 4)\t0\n", - "[4, 6)\t0\n", - "[6, 8)\t0\n", - "[8, 10)\t0\n", - "[10, inf) 0\n", - "7568 total\n", - "\n", - "social_rules_rewards\n", - "bucket\tcount\n", - "(-inf, 0)\t 44\n", - "[0, 2)\t7524\n", - "[2, 4)\t0\n", - "[4, 6)\t0\n", - "[6, 8)\t0\n", - "[8, 10)\t0\n", - "[10, inf) 0\n", - "7568 total\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "from collections import Counter\n", "plot_6_buckets(Counter(goal_rewards), \"goal_0503_rewards\")\n", @@ -297,7 +171,7 @@ }, { "cell_type": "code", - "execution_count": 39, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -313,7 +187,7 @@ }, { "cell_type": "code", - "execution_count": 40, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -332,7 +206,7 @@ }, { "cell_type": "code", - "execution_count": 41, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -356,7 +230,7 @@ }, { "cell_type": "code", - "execution_count": 42, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -371,7 +245,7 @@ }, { "cell_type": "code", - "execution_count": 43, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -386,27 +260,15 @@ " \"knowledge\": normalized_knowledge_values,\n", " \"goal_relationship_knowledge\": normalized_goal_relationship_knowledge_map_values,\n", "}\n", - "import pickle\n", "with open(\"reward_value_dicts.json\", 'w') as f:\n", " json.dump(value_dict, f)" ] }, { "cell_type": "code", - "execution_count": 44, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Key: goal, Type: , Length: 7568\n", - "Key: relationship, Type: , Length: 7568\n", - "Key: knowledge, Type: , Length: 7568\n", - "Key: goal_relationship_knowledge, Type: , Length: 7568\n" - ] - } - ], + "outputs": [], "source": [ "for key, value in value_dict.items():\n", " print(f\"Key: {key}, Type: {type(value)}, Length: {len(value)}\")" @@ -414,7 +276,7 @@ }, { "cell_type": "code", - "execution_count": 45, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -441,30 +303,9 @@ }, { "cell_type": "code", - "execution_count": 46, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Correlation between goal and overall:\n", - "(0.8687566540029152, 0.0)\n", - "Correlation between goal and financial:\n", - "(0.1401024952037344, 1.7531092600883863e-34)\n", - "Correlation between goal and knowledge:\n", - "(0.47331812212986424, 0.0)\n", - "Correlation between goal and relationship:\n", - "(0.4594666455638434, 0.0)\n", - "Correlation between goal and believability:\n", - "(0.6710441213034302, 0.0)\n", - "Correlation between goal and secret:\n", - "(-0.01704789687807092, 0.13809252367705063)\n", - "Correlation between goal and social rules:\n", - "(-0.04061920984921264, 0.00040852772938770995)\n" - ] - } - ], + "outputs": [], "source": [ "print(\"Correlation between goal and overall:\")\n", "print(calculate_correlation(goal_map, overall_map))\n", @@ -484,19 +325,9 @@ }, { "cell_type": "code", - "execution_count": 47, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.8447560493459538, 0.0)\n", - "1.5204222457129433\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -529,19 +360,9 @@ }, { "cell_type": "code", - "execution_count": 48, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.8783043496324526, 0.0)\n", - "1.2408386187455958\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -570,19 +391,9 @@ }, { "cell_type": "code", - "execution_count": 49, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.8783043496324526, 0.0)\n", - "1.2408386187455958\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -611,19 +422,9 @@ }, { "cell_type": "code", - "execution_count": 50, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.8991990776133605, 0.0)\n", - "1.4942080690627202\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -652,19 +453,9 @@ }, { "cell_type": "code", - "execution_count": 51, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.5591519503039457, 0.0)\n", - "2.2806333685694153\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -693,19 +484,9 @@ }, { "cell_type": "code", - "execution_count": 52, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(1.0, 0.0)\n", - "0.0\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", @@ -730,19 +511,9 @@ }, { "cell_type": "code", - "execution_count": 54, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "7568\n", - "(0.45946664556384337, 0.0)\n", - "2.8886539816772374\n" - ] - } - ], + "outputs": [], "source": [ "new_reward = []\n", "for entry in goal:\n", diff --git a/scripts/annotate/post_process_annotation_rm.py b/scripts/annotate/post_process_annotation_rm.py index 91028d9..59b1652 100644 --- a/scripts/annotate/post_process_annotation_rm.py +++ b/scripts/annotate/post_process_annotation_rm.py @@ -50,14 +50,14 @@ def get_attributed_data(data: List[Dict[str, Any]], utterance_pattern: str, do_n new_utterance['goal_score'] = d['goal_score'] total_attributed_reward += new_utterance['attributed_reward'] all_utterances.append(new_utterance) - + if do_normalize: for utterance in all_utterances: utterance['attributed_reward'] = utterance['attributed_reward'] / total_attributed_reward * d['goal_score'] - + for utterance in all_utterances: attributed_data.append(utterance) - + print(f"Error Count: {error_count}") return attributed_data @@ -80,7 +80,7 @@ def main(data_dir: str, input_file: str, reward_output_file: str, do_normalize: utterance_pattern = r'Utterance (\d+) by ([A-Za-z ]+)' print("turning into attributed utterances") - + attributed_data = get_attributed_data(data, utterance_pattern, do_normalize=do_normalize) sotopia_pi_utterance_reward = [] # breakpoint() diff --git a/scripts/data_process/data_filtering.py b/scripts/data_process/data_filtering.py index 87205ec..2e7c5cc 100644 --- a/scripts/data_process/data_filtering.py +++ b/scripts/data_process/data_filtering.py @@ -5,12 +5,11 @@ import json import random -from collections import Counter -from sotopia.database import episodes_to_jsonl import matplotlib.pyplot as plt import numpy as np from prompt_reverse_engineering import reverse_episode_log +from sotopia.database import episodes_to_jsonl from sotopia.database.logs import EpisodeLog @@ -304,7 +303,7 @@ def behavior_cloning_data_filtering_mistral_mistral(): eps_list += eps_by_env_filtered[env] print(total_num) - + def behavior_cloning_data_filtering_qwen_qwen(): # get conversation episodes according to social task scenarios eps_list = [] @@ -320,7 +319,7 @@ def behavior_cloning_data_filtering_qwen_qwen(): print(len(eps_by_env)) eps_by_env_filtered = filter_episodes_on_top_2_selfplay(eps_by_env) - + total_num = 0 for env in eps_by_env_filtered: total_num += len(eps_by_env_filtered[env]) @@ -334,4 +333,4 @@ def behavior_cloning_data_filtering_qwen_qwen(): if __name__ == "__main__": # behavior_cloning_data_filtering_gpt4_gpt4() # behavior_cloning_data_filtering_mistral_mistral() - behavior_cloning_data_filtering_qwen_qwen() \ No newline at end of file + behavior_cloning_data_filtering_qwen_qwen() diff --git a/scripts/data_process/db_free_data_filtering.py b/scripts/data_process/db_free_data_filtering.py index ba86011..7e85dda 100644 --- a/scripts/data_process/db_free_data_filtering.py +++ b/scripts/data_process/db_free_data_filtering.py @@ -1,11 +1,9 @@ -import os # TODO: Fill in REDIS OM URL in the form of `redis://:password@host:port` # os.environ["REDIS_OM_URL"] = "redis://:password@server_name:port_num" import json import random -from collections import Counter import matplotlib.pyplot as plt import numpy as np @@ -312,4 +310,4 @@ def behavior_cloning_data_filtering_mistral_mistral(): if __name__ == "__main__": behavior_cloning_data_filtering_gpt4_gpt4() - # behavior_cloning_data_filtering_mistral_mistral() \ No newline at end of file + # behavior_cloning_data_filtering_mistral_mistral() diff --git a/scripts/data_process/generate_conversations.py b/scripts/data_process/generate_conversations.py index fd38d85..8783072 100644 --- a/scripts/data_process/generate_conversations.py +++ b/scripts/data_process/generate_conversations.py @@ -1,7 +1,6 @@ import argparse import json import re -import subprocess def overwrite_eval_bash( @@ -123,4 +122,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/data_process/generate_sft_from_episodes.py b/scripts/data_process/generate_sft_from_episodes.py index cbb7121..ecf8215 100644 --- a/scripts/data_process/generate_sft_from_episodes.py +++ b/scripts/data_process/generate_sft_from_episodes.py @@ -1,10 +1,12 @@ -import click +import glob import json import os -import glob from typing import Any, Dict, List -from tqdm import tqdm + +import click from db_free_reverse_engineering import run_reverse_by_pk_agent +from tqdm import tqdm + @click.command() @click.option("--data_dir", type=str, required=True, help="Directory containing data files.") @@ -15,30 +17,30 @@ def main(data_dir: str, utterances_output_subdir: str, episodes_file: str, sft_o episode_path = os.path.join(data_dir, episodes_file) if not os.path.exists(episode_path): raise Exception(f"Episodes file not found: {episode_path}") - + with open(episode_path, 'r') as f: data: List[Dict[str, Any]] = [json.loads(d) for d in f.readlines()] - + cache_dir = os.path.join(data_dir, utterances_output_subdir) if not os.path.exists(cache_dir): os.makedirs(cache_dir) for d in tqdm(data): run_reverse_by_pk_agent(d['episode_id'], True, cache_dir, episode_path) run_reverse_by_pk_agent(d['episode_id'], False, cache_dir, episode_path) - + utterances = [] for record in glob.glob(f"{cache_dir}/*.json"): with open(record, 'r') as f: uttr = json.load(f) utterances.append(uttr) - + sft_utterances = [] for uttr in utterances: sft_utterances.append({ "input": uttr['prompt'] + " Your available action types are\nspeak none action leave non-verbal communication.\nNote: 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\nPlease only generate a JSON string including the action type and the argument.\nYour action should follow the given format:\nThe output should be formatted as a JSON instance that conforms to the JSON schema below.\n\nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}\nthe object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\n\nHere is the output schema:\n```\n{\"properties\": {\"action_type\": {\"description\": \"whether to speak at this turn or choose to not do anything\", \"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"title\": \"Action Type\", \"type\": \"string\"}, \"argument\": {\"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", \"title\": \"Argument\", \"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}\n```", "output": uttr['result'], }) - + print(f"Total utterances: {len(sft_utterances)}") with open(os.path.join(data_dir, sft_output_file), 'w') as f: json.dump(sft_utterances, f, indent=4) diff --git a/scripts/data_process/serialize.py b/scripts/data_process/serialize.py index 4ffe740..0b4ee22 100644 --- a/scripts/data_process/serialize.py +++ b/scripts/data_process/serialize.py @@ -1,5 +1,5 @@ -from sotopia.database import episodes_to_jsonl, EpisodeLog - +from sotopia.database import EpisodeLog, episodes_to_jsonl + episodes: list[EpisodeLog] = EpisodeLog.find(EpisodeLog.tag=="Qwen2.5-7b-Instruct_vs_Qwen2.5-7b-Instruct-0510") - -episodes_to_jsonl(episodes, "Qwen2.5-7b-Instruct_vs_Qwen2.5-7b-Instruct-0510.jsonl") \ No newline at end of file + +episodes_to_jsonl(episodes, "Qwen2.5-7b-Instruct_vs_Qwen2.5-7b-Instruct-0510.jsonl") diff --git a/scripts/data_process/sotopia_all_eval_script.sh b/scripts/data_process/sotopia_all_eval_script.sh index b26aae3..7981f2f 100644 --- a/scripts/data_process/sotopia_all_eval_script.sh +++ b/scripts/data_process/sotopia_all_eval_script.sh @@ -24,4 +24,4 @@ python examples/experiment_eval.py \ '--gin.BATCH_SIZE=20' \ '--gin.TAG="xx"' \ '--gin.PUSH_TO_DB=True' \ -'--gin.TAG_TO_CHECK_EXISTING_EPISODES="qwen-sft-qwen-sft-3-26-v2"' \ No newline at end of file +'--gin.TAG_TO_CHECK_EXISTING_EPISODES="qwen-sft-qwen-sft-3-26-v2"' diff --git a/scripts/data_process/sotopia_pi_and_eval.sh b/scripts/data_process/sotopia_pi_and_eval.sh index ef0d868..e059131 100644 --- a/scripts/data_process/sotopia_pi_and_eval.sh +++ b/scripts/data_process/sotopia_pi_and_eval.sh @@ -51,4 +51,4 @@ python generate_sft_from_episodes.py \ --data_dir ../../data \ --utterances_output_subdir sotopia_pi_round1_qwen_utterances_filtered \ --episodes_file sotopia_pi_round1_qwen_episodes_filtered.jsonl \ ---sft_output_file sotopia_pi_round1_qwen_sft_pi.json \ No newline at end of file +--sft_output_file sotopia_pi_round1_qwen_sft_pi.json diff --git a/scripts/data_process/sotopia_pi_self_play_script.sh b/scripts/data_process/sotopia_pi_self_play_script.sh index d922c37..68b7172 100644 --- a/scripts/data_process/sotopia_pi_self_play_script.sh +++ b/scripts/data_process/sotopia_pi_self_play_script.sh @@ -43,4 +43,4 @@ python examples/experiment_eval.py \ '--gin.BATCH_SIZE=20' \ '--gin.TAG="qwen-sft-qwen-sft-3-26-v2"' \ '--gin.PUSH_TO_DB=True' \ -'--gin.TAG_TO_CHECK_EXISTING_EPISODES="qwen-sft-qwen-sft-3-26-v2"' \ No newline at end of file +'--gin.TAG_TO_CHECK_EXISTING_EPISODES="qwen-sft-qwen-sft-3-26-v2"' diff --git a/scripts/data_process/used_env.json b/scripts/data_process/used_env.json index 9168e9a..c9bd67e 100644 --- a/scripts/data_process/used_env.json +++ b/scripts/data_process/used_env.json @@ -527,4 +527,4 @@ "01H7VFHP8AN5643B0NR0NP00VE", "01H7VFHN7A1ZX5KSMT2YN9RXC4" ] -} \ No newline at end of file +} diff --git a/scripts/inference_rm.py b/scripts/inference_rm.py index 09e81dd..377b8b4 100644 --- a/scripts/inference_rm.py +++ b/scripts/inference_rm.py @@ -96,4 +96,4 @@ def main(): print("GTH REWARD: Not available") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/train_grpo.py b/scripts/train_grpo.py index f4ed487..849d96f 100644 --- a/scripts/train_grpo.py +++ b/scripts/train_grpo.py @@ -1,13 +1,14 @@ import argparse import os + os.environ["TRANSFORMERS_NO_COMPILE"] = "1" -import argparse -from sotopia_rl import SotopiaGRPOTrainer from accelerate import Accelerator +from sotopia_rl import SotopiaGRPOTrainer + if __name__ == '__main__': parser = argparse.ArgumentParser(description="Train a model with GRPO using a reward model.") - + parser.add_argument("--model_name", type=str, default="/data/models/gemma-2-2b-it", help="Path to the model") @@ -57,7 +58,7 @@ help="Wandb run name") parser.add_argument("--use_lora_train_grpo", action="store_true", help="Use LoRA for training GRPO") - + args = parser.parse_args() accelerator = Accelerator() trainer = SotopiaGRPOTrainer(args, accelerator) diff --git a/scripts/train_grpo.sh b/scripts/train_grpo.sh index ddbb273..9a8fd31 100644 --- a/scripts/train_grpo.sh +++ b/scripts/train_grpo.sh @@ -15,4 +15,4 @@ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 accelerate launch \ --num_grpo_epochs 2 \ --use_lora_train_grpo \ --num_generations 16 \ - --output_dir ../grpo_checkpoints_qwen2.5-7b + --output_dir ../grpo_checkpoints_qwen2.5-7b diff --git a/scripts/train_rm.py b/scripts/train_rm.py index 236a8c8..3a64296 100644 --- a/scripts/train_rm.py +++ b/scripts/train_rm.py @@ -1,6 +1,7 @@ import argparse -import os + from accelerate import Accelerator + from sotopia_rl import SotopiaRMTrainer if __name__ == '__main__': diff --git a/scripts/train_sft.py b/scripts/train_sft.py index 48fceae..fe3df95 100644 --- a/scripts/train_sft.py +++ b/scripts/train_sft.py @@ -1,6 +1,7 @@ import argparse -import os + from accelerate import Accelerator + from sotopia_rl import SotopiaSFTTrainer if __name__ == "__main__": @@ -34,4 +35,4 @@ accelerator = Accelerator() trainer = SotopiaSFTTrainer(args, accelerator) - trainer.train() \ No newline at end of file + trainer.train() diff --git a/sotopia_rl/__init__.py b/sotopia_rl/__init__.py index d90bb5b..799c8be 100644 --- a/sotopia_rl/__init__.py +++ b/sotopia_rl/__init__.py @@ -1,6 +1,6 @@ +from .grpo_trainer import SotopiaGRPOTrainer +from .ppo_trainer import SotopiaPPOTrainer from .rm_trainer import SotopiaRMTrainer from .sft_trainer import SotopiaSFTTrainer -from .ppo_trainer import SotopiaPPOTrainer -from .grpo_trainer import SotopiaGRPOTrainer __all__ = ["SotopiaPPOTrainer", "SotopiaRMTrainer", "SotopiaSFTTrainer", "SotopiaGRPOTrainer"] diff --git a/sotopia_rl/data.py b/sotopia_rl/data.py index b7fff53..c15c9b8 100644 --- a/sotopia_rl/data.py +++ b/sotopia_rl/data.py @@ -172,11 +172,11 @@ def collate_fn(self, batch): [item["input_ids"] for item in batch], batch_first=True, padding_value=self.tokenizer.pad_token_id ) return {"input_ids": input_ids} - + class GRPODataset(Dataset): def __init__(self, data_path: str, tokenizer, template, max_length: int): self.data = self.load_sft_data(data_path) - self.tokenizer = tokenizer + self.tokenizer = tokenizer self.max_length = max_length self.template = template @@ -197,5 +197,5 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: return { "prompt": rendered_prompt, - "completion": item["output"] + "completion": item["output"] } diff --git a/sotopia_rl/grpo_trainer.py b/sotopia_rl/grpo_trainer.py index b520715..4d7652a 100644 --- a/sotopia_rl/grpo_trainer.py +++ b/sotopia_rl/grpo_trainer.py @@ -1,24 +1,20 @@ import os + import torch import wandb -from datasets import load_dataset +from accelerate import Accelerator +from jinja2 import Environment, FileSystemLoader +from peft import PeftModelForCausalLM, PeftModelForSequenceClassification from torch.utils.data import random_split from transformers import ( AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, - GenerationConfig, ) -from accelerate import PartialState -from peft import PeftModelForCausalLM, PeftModelForSequenceClassification -from jinja2 import Environment, FileSystemLoader -from trl import get_kbit_device_map, GRPOConfig, GRPOTrainer -from accelerate import Accelerator +from trl import GRPOConfig, GRPOTrainer, get_kbit_device_map + from sotopia_rl.data import GRPODataset -from functools import partial -from typing import List -from peft import prepare_model_for_kbit_training os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" os.environ["NCCL_P2P_DISABLE"] = "1" diff --git a/sotopia_rl/ppo_trainer.py b/sotopia_rl/ppo_trainer.py index bde99fc..433e45f 100644 --- a/sotopia_rl/ppo_trainer.py +++ b/sotopia_rl/ppo_trainer.py @@ -1,6 +1,8 @@ import os import torch +import wandb +from accelerate import Accelerator, PartialState from jinja2 import Environment, FileSystemLoader from peft import PeftModelForCausalLM, PeftModelForSequenceClassification from torch.utils.data import random_split @@ -9,14 +11,12 @@ AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, - GenerationConfig, ) -from trl import get_kbit_device_map, PPOConfig, PPOTrainer -from accelerate import PartialState, Accelerator +from trl import PPOConfig, PPOTrainer, get_kbit_device_map -import wandb from sotopia_rl.data import PPODataset -os.environ['NCCL_P2P_DISABLE'] = '1' + +os.environ['NCCL_P2P_DISABLE'] = '1' os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL" class SotopiaPPOTrainer: @@ -81,7 +81,7 @@ def _setup_dataset(self): max_length=self.args.max_length ) print(f"dataset: {len(dataset)}") - + generator = torch.Generator().manual_seed(42) val_ratio = getattr(self.args, 'val_ratio', 0.05) train_size = min(int(len(dataset) * (1 - val_ratio)), len(dataset) - 2) @@ -96,7 +96,7 @@ def _create_quantization_config(self): bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4" ) - + def _setup_generation_models(self): base_gen_ref = AutoModelForCausalLM.from_pretrained( self.args.model_name, @@ -181,14 +181,14 @@ def _setup_classification_models(self): torch_dtype='auto', num_labels=1, ) - + # VERY VERY IMPORTANT - # specifically designed for PPO training, + # specifically designed for PPO training, # based on the get_reward function # it fill the input_ids paddings with 0s self.value_model.config.pad_token_id = 0 self.reward_model.config.pad_token_id = 0 - + requires_grad_num = 0 for name, param in self.value_model.named_parameters(): if param.requires_grad: @@ -240,4 +240,4 @@ def train(self): return train_stats except Exception as e: print(f"Training error: {str(e)}") - raise \ No newline at end of file + raise diff --git a/sotopia_rl/prompter/attribution_methods.py b/sotopia_rl/prompter/attribution_methods.py index b789306..9ec7f74 100644 --- a/sotopia_rl/prompter/attribution_methods.py +++ b/sotopia_rl/prompter/attribution_methods.py @@ -1,6 +1,12 @@ +from sotopia_rl.prompter.all_the_same_attribution_function import ( + get_attribution_single_conv as all_the_same_attribution_normalized_single_conv, +) from sotopia_rl.prompter.direct_attribution_function import ( get_attribution_single_conv as direct_attribution_single_conv, ) +from sotopia_rl.prompter.direct_attribution_generic_function import ( + get_attribution_single_conv as direct_attribution_generic_single_conv, +) from sotopia_rl.prompter.direct_attribution_normalized_function import ( get_attribution_single_conv as direct_attribution_normalized_single_conv, ) @@ -10,6 +16,9 @@ from sotopia_rl.prompter.goal_progress_attribution_function import ( get_attribution_single_conv as goal_progress_attribution_single_conv, ) +from sotopia_rl.prompter.key_utterance_function import ( + get_attribution_single_conv as key_utterance_attribution_single_conv, +) from sotopia_rl.prompter.only_response_attribution_function import ( get_attribution_single_conv as only_response_attribution_single_conv, ) @@ -19,15 +28,7 @@ from sotopia_rl.prompter.utterance_quality_attribution_normalized_function import ( get_attribution_single_conv as utterance_quality_attribution_normalized_single_conv, ) -from sotopia_rl.prompter.all_the_same_attribution_function import ( - get_attribution_single_conv as all_the_same_attribution_normalized_single_conv, -) -from sotopia_rl.prompter.key_utterance_function import ( - get_attribution_single_conv as key_utterance_attribution_single_conv, -) -from sotopia_rl.prompter.direct_attribution_generic_function import ( - get_attribution_single_conv as direct_attribution_generic_single_conv, -) + # from sotopia_rl.prompter.utterance_quality_generic_function import ( # get_attribution_single_conv as utterance_quality_attribution_generic_single_conv, # ) @@ -44,4 +45,4 @@ "key_utterance": key_utterance_attribution_single_conv, "direct_generic": direct_attribution_generic_single_conv, # "utterance_quality_generic": utterance_quality_attribution_generic_single_conv, -} \ No newline at end of file +} diff --git a/sotopia_rl/prompter/direct_attribution_function.py b/sotopia_rl/prompter/direct_attribution_function.py index 9b3d42e..f61d6b0 100644 --- a/sotopia_rl/prompter/direct_attribution_function.py +++ b/sotopia_rl/prompter/direct_attribution_function.py @@ -4,9 +4,7 @@ from openai import OpenAI -from sotopia_rl.prompter.one_pass_instructions import ( - ATTRIBUTION_INSTRUCTIONS_DICT, -) +from sotopia_rl.prompter.one_pass_instructions import ATTRIBUTION_INSTRUCTIONS_DICT def openai_call(prompt: str, model: str = "gpt-3.5-turbo") -> str | None: @@ -85,7 +83,7 @@ def assign_attributions_for_conversation( except ValueError: print("Failed to convert all values to integers; retrying") continue - + if uttr_count != len(result) and i < 4: print("Response length does not match the number of agent utterances; retrying") elif uttr_count == len(result): diff --git a/sotopia_rl/prompter/direct_attribution_generic_function.py b/sotopia_rl/prompter/direct_attribution_generic_function.py index 52afc69..623ff5f 100644 --- a/sotopia_rl/prompter/direct_attribution_generic_function.py +++ b/sotopia_rl/prompter/direct_attribution_generic_function.py @@ -1,13 +1,17 @@ import json +import os import re from typing import Any, Dict, List, Tuple -import os + from openai import OpenAI from sotopia_rl.prompter.generic_templates import ( - SCALE_GUIDELINE_DICT, DIMENSION_DESCRIPTION_DICT, DIRECT_ATTRIBUTION_TEMPLATE, + DIMENSION_DESCRIPTION_DICT, + DIRECT_ATTRIBUTION_TEMPLATE, + SCALE_GUIDELINE_DICT, ) + def openai_call(prompt: str, model: str = "gpt-3.5-turbo") -> str | None: if model in ["gpt-3.5-turbo", "gpt-4", "gpt-4o", "o4-mini"]: client = OpenAI() @@ -135,7 +139,7 @@ def assign_attributions_for_conversation( except ValueError: print("Failed to convert all values to integers; retrying") continue - + if uttr_count != len(result) and i < 4: print("Response length does not match the number of agent utterances; retrying") elif uttr_count == len(result): @@ -189,7 +193,7 @@ def get_attribution_single_conv(conversation, agent, goals, episode, rewards, ll dim_score = 10 elif dimension == "goal_barebone": dim_score = rewards[agent]["goal"] - else: + else: dim_score = rewards[agent][dimension] attribution_rewards = calc_attributed_reward(attribution_scores, scale, dim_score) for key in attribution_rewards: diff --git a/sotopia_rl/prompter/direct_attribution_normalized_function.py b/sotopia_rl/prompter/direct_attribution_normalized_function.py index f737c1e..6f9dc4c 100644 --- a/sotopia_rl/prompter/direct_attribution_normalized_function.py +++ b/sotopia_rl/prompter/direct_attribution_normalized_function.py @@ -4,9 +4,7 @@ from openai import OpenAI -from sotopia_rl.prompter.one_pass_instructions import ( - ATTRIBUTION_INSTRUCTIONS_DICT, -) +from sotopia_rl.prompter.one_pass_instructions import ATTRIBUTION_INSTRUCTIONS_DICT def openai_call(prompt: str, model: str = "gpt-3.5-turbo") -> str | None: @@ -79,7 +77,7 @@ def assign_attributions_for_conversation( print(response) return {} result = json.loads(formatted_response) - + if uttr_count != len(result) and i < 4: print("Response length does not match the number of agent utterances; retrying") elif uttr_count == len(result): @@ -100,7 +98,7 @@ def calc_attributed_reward(attributed_data: List[Dict[str, float | int]], attrib for k, v in attributed_data.items(): total_attributions += v for k, v in attributed_data.items(): - utterance_reward_map[k] = {"reward": calc_reward(v, attribution_instruction_name, goal_score, total_attributions), + utterance_reward_map[k] = {"reward": calc_reward(v, attribution_instruction_name, goal_score, total_attributions), "attribution": v} return utterance_reward_map diff --git a/sotopia_rl/prompter/generic_templates.py b/sotopia_rl/prompter/generic_templates.py index 100163d..57b1b42 100644 --- a/sotopia_rl/prompter/generic_templates.py +++ b/sotopia_rl/prompter/generic_templates.py @@ -71,7 +71,7 @@ - Note that you should only consider the contribution to the goal achievement. For each utterance, assess whether the goal is achieved. If a goal is already achieved, the utterance should not be assigned a score higher than 1. """ -CONVERSATION_BEHAVIOR_DESCRIPTION = """ +CONVERSATION_BEHAVIOR_DESCRIPTION = """ Conversation Behavior refers to the evaluation of the agent's conversation behavior, including the avoidance of repetitive utterances, proper ending of the conversation, and overall performance as a social agent. A higher score indicates effective conversation behavior, while a lower score indicates poor performance in these areas. DOMAIN SPECIFIC SCORING GUIDELINES: @@ -138,4 +138,4 @@ 10. Formatting Instructions: {formatting_instructions} -""" \ No newline at end of file +""" diff --git a/sotopia_rl/prompter/key_utterance_function.py b/sotopia_rl/prompter/key_utterance_function.py index 0aca7fb..47d0176 100644 --- a/sotopia_rl/prompter/key_utterance_function.py +++ b/sotopia_rl/prompter/key_utterance_function.py @@ -1,12 +1,9 @@ -import json import re from typing import Any, Dict, List, Tuple from openai import OpenAI -from sotopia_rl.prompter.one_pass_instructions import ( - ATTRIBUTION_INSTRUCTIONS_DICT, -) +from sotopia_rl.prompter.one_pass_instructions import ATTRIBUTION_INSTRUCTIONS_DICT REGEX = "^Utterance (?:[0-9]|[1-9][0-9]) by {agent}$" @@ -62,7 +59,7 @@ def assign_attributions_for_conversation( uttr_count += 1 uttr_attr_dict[f"Utterance {j//2} by {speaker}"] = 0 response = openai_call(prompt, llm_name).strip() - + if response is None: print("Failed to get response from OpenAI; returning empty dictionary") return {} diff --git a/sotopia_rl/prompter/one_pass_instructions.py b/sotopia_rl/prompter/one_pass_instructions.py index 56f791a..15e5513 100644 --- a/sotopia_rl/prompter/one_pass_instructions.py +++ b/sotopia_rl/prompter/one_pass_instructions.py @@ -63,4 +63,4 @@ "default": DEFAULT_DIRECT_INSTRUCTIONS, "10-scale": DIRECT_10_SCALE_INSTRUCTIONS, "key_utterance": KEY_UTTERANCE_INSTRUCTIONS, -} \ No newline at end of file +} diff --git a/sotopia_rl/prompter/only_response_attribution_function.py b/sotopia_rl/prompter/only_response_attribution_function.py index 9954e7d..1cf735b 100644 --- a/sotopia_rl/prompter/only_response_attribution_function.py +++ b/sotopia_rl/prompter/only_response_attribution_function.py @@ -26,7 +26,7 @@ ONLY_RESPONSE_DIRECT_INSTRUCTIONS = """ ## Reward Attribution Instructions for LLMs -Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. +Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. 1. Input Context: - You will receive the utterance or action of an agent in one turn of the conversation. @@ -69,7 +69,7 @@ DEFAULT_PROMPT = """ ## Reward Attribution Instructions for LLMs -Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. +Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. 1. Input Context: - You will receive the utterance or action of an agent in one turn of the conversation and the conversation before it. diff --git a/sotopia_rl/prompter/utterance_quality_attribution_function.py b/sotopia_rl/prompter/utterance_quality_attribution_function.py index 06e0503..c055478 100644 --- a/sotopia_rl/prompter/utterance_quality_attribution_function.py +++ b/sotopia_rl/prompter/utterance_quality_attribution_function.py @@ -10,7 +10,7 @@ DEFAULT_PROMPT = """ ## Reward Attribution Instructions for LLMs -Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. +Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. 1. Input Context: - You will recieve the utterance or action of an agent at a certain point and the conversation before it. diff --git a/sotopia_rl/prompter/utterance_quality_attribution_normalized_function.py b/sotopia_rl/prompter/utterance_quality_attribution_normalized_function.py index ad8be3e..b9627d2 100644 --- a/sotopia_rl/prompter/utterance_quality_attribution_normalized_function.py +++ b/sotopia_rl/prompter/utterance_quality_attribution_normalized_function.py @@ -26,7 +26,7 @@ ONLY_RESPONSE_DIRECT_INSTRUCTIONS = """ ## Reward Attribution Instructions for LLMs -Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. +Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. 1. Input Context: - You will recieve the utterance or action of an agent at a certain point and the conversation before it. @@ -66,7 +66,7 @@ DEFAULT_PROMPT = """ ## Reward Attribution Instructions for LLMs -Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. +Two agents are in a conversation. For now, you are the judge of the utterance of one of the agents. 1. Input Context: - You will recieve the utterance or action of an agent at a certain point and the conversation before it. diff --git a/sotopia_rl/rm_trainer.py b/sotopia_rl/rm_trainer.py index 47acbd5..776f9bd 100644 --- a/sotopia_rl/rm_trainer.py +++ b/sotopia_rl/rm_trainer.py @@ -1,9 +1,10 @@ import os import torch -import torch.distributed as dist +import torch._dynamo +import wandb from jinja2 import Environment, FileSystemLoader -from peft import LoraConfig, get_peft_model, PeftModelForSequenceClassification +from peft import LoraConfig, PeftModelForSequenceClassification from torch.nn import MSELoss from torch.utils.data import random_split from transformers import ( @@ -12,12 +13,8 @@ Trainer, TrainingArguments, ) -from accelerate import Accelerator -from typing import Optional -import wandb from sotopia_rl.data import RMDataset -import torch._dynamo torch._dynamo.config.suppress_errors = True diff --git a/sotopia_rl/sft_trainer.py b/sotopia_rl/sft_trainer.py index 41dcca0..a7d2d63 100644 --- a/sotopia_rl/sft_trainer.py +++ b/sotopia_rl/sft_trainer.py @@ -1,23 +1,20 @@ import os +from functools import partial import torch -from functools import partial -from torch.nn.utils.rnn import pad_sequence +import wandb from jinja2 import Environment, FileSystemLoader -from peft import LoraConfig, get_peft_model, PeftModelForCausalLM -from torch.utils.data import random_split +from torch.nn.utils.rnn import pad_sequence from transformers import ( - Trainer, AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, + Trainer, TrainingArguments, ) -from trl import SFTTrainer -import wandb + from sotopia_rl.data import SFTDataset -from datasets import Dataset os.environ['NCCL_P2P_DISABLE'] = '1' os.environ["TOKENIZERS_PARALLELISM"] = "false" diff --git a/sotopia_rl/utils/openai.py b/sotopia_rl/utils/openai.py index f89ab57..b175627 100644 --- a/sotopia_rl/utils/openai.py +++ b/sotopia_rl/utils/openai.py @@ -1,6 +1,8 @@ -from openai import OpenAI import os +from openai import OpenAI + + def openai_call(prompt: str, model: str = "gpt-3.5-turbo") -> str | None: if model in ["gpt-3.5-turbo", "gpt-4", "gpt-4o", "o4-mini"]: client = OpenAI() @@ -52,4 +54,4 @@ def openai_call(prompt: str, model: str = "gpt-3.5-turbo") -> str | None: ) return response.choices[0].text else: - raise ValueError(f"Model {model} not supported.") \ No newline at end of file + raise ValueError(f"Model {model} not supported.")