Update training.py

training.py

@@ -1,5 +1,4 @@
-# training.py – FIXED PPO training (no variable names changed)
-
+# training.py 
 import json
 import torch
 import torch.nn.functional as F
@@ -20,7 +19,7 @@ from environment import CodeReviewEnv
 from models import (
     RunTests, RunLinter, Inspect,
     ProposeFix, WriteComment, AskQuestion,
-    Done, Skip
+    Done, Skip , QueryDocs 
 )
 
 # ======================================================================
@@ -85,6 +84,8 @@ def map_to_env(action: AgentAction):
         return WriteComment(comment_text=action.content or "")
     elif action.action_type == "question":
         return AskQuestion(question=action.content or "")
+    elif action.action_type == "query_docs":               # <-- new
+        return QueryDocs(query_topic=action.content or "")
     elif action.action_type == "done":
         return Done()
     else: