fix the format

Author: qywu

rllm/engine/agent_workflow_engine.py

@@ -49,7 +49,7 @@ def __init__(self, workflow_cls: type[Workflow], workflow_args: dict, rollout_en
         self.n_parallel_tasks = n_parallel_tasks
         self.executor = ThreadPoolExecutor(max_workers=self.n_parallel_tasks)
         self.workflow_queue = None
-        
+
         # Episode logging support
         self.episode_logger = episode_logger
         self.current_step = 0
@@ -58,7 +58,7 @@ def __init__(self, workflow_cls: type[Workflow], workflow_args: dict, rollout_en
 
     def set_training_step(self, step: int, mode: str = "train", epoch: int = 0):
         """Set current training step for episode logging.
-        
+
         Args:
             step: Current training step number
             mode: Mode identifier ('train' or 'val'), defaults to 'train'
@@ -67,7 +67,7 @@ def set_training_step(self, step: int, mode: str = "train", epoch: int = 0):
         self.current_step = step
         self.current_mode = mode
         self.current_epoch = epoch
-    
+
     async def initialize_pool(self):
         """Initialize the workflow pool with parallel workflow instances.
 
@@ -173,7 +173,7 @@ async def execute_tasks(self, tasks: list[dict], task_ids: list[str] | None = No
         sorted_tasks = sorted(task_states.keys(), key=lambda task_id: task_states[task_id]["idx"])
         for task_id in sorted_tasks:
             results.extend(task_states[task_id]["episodes"])
-        
+
         # Log episodes if logger is provided
         if self.episode_logger is not None:
             try:
@@ -182,8 +182,9 @@ async def execute_tasks(self, tasks: list[dict], task_ids: list[str] | None = No
             except Exception as e:
                 logger.error(f"Failed to log episodes: {e}")
                 import traceback
+
                 traceback.print_exc()
-        
+
         return results
 
     async def execute_tasks_verl(self, batch: "DataProto", **kwargs) -> "DataProto":

rllm/trainer/verl/agent_workflow_trainer.py

@@ -79,13 +79,10 @@ def init_workers(self):
 
         # Create episode logger if enabled in config
         episode_logger = None
-        if self.config.trainer.get('log_episodes', False):
+        if self.config.trainer.get("log_episodes", False):
             # Get episode log directory from config, default to "logs/my_project/my_experiment"
-            episode_log_dir = self.config.trainer.get('episode_log_dir', f'logs/{self.config.trainer.project_name}/{self.config.trainer.experiment_name}')
-            episode_logger = EpisodeLogger(
-                base_dir=episode_log_dir,
-                subdirectory="episodes"
-            )
+            episode_log_dir = self.config.trainer.get("episode_log_dir", f"logs/{self.config.trainer.project_name}/{self.config.trainer.experiment_name}")
+            episode_logger = EpisodeLogger(base_dir=episode_log_dir, subdirectory="episodes")
 
         self.agent_execution_engine = AgentWorkflowEngine(
             workflow_cls=self.workflow_class,

rllm/utils/__init__.py

@@ -3,4 +3,3 @@
 from rllm.utils.episode_logger import EpisodeLogger
 
 __all__ = ["EpisodeLogger"]
-

rllm/utils/episode_logger.py

@@ -1,143 +1,135 @@
 """Episode JSON Logger for saving detailed episode information."""
-import json
+
 import hashlib
+import json
 from pathlib import Path
 from typing import Any
+
 from rllm.agents.agent import Episode
 
 
 class EpisodeLogger:
     """Logger to save episodes to individual JSON files with step and data hash."""
-    
+
     def __init__(self, base_dir: str, subdirectory: str = "episodes"):
         """Initialize the episode logger.
-        
+
         Args:
-            base_dir: Base directory for episode logs. Can be configured via 
-                     config.trainer.episode_log_dir 
+            base_dir: Base directory for episode logs. Can be configured via
+                     config.trainer.episode_log_dir
                      (default: "logs/${trainer.project_name}/${trainer.experiment_name}")
             subdirectory: Subdirectory within base_dir for episodes (default: "episodes")
                          Final path will be: {base_dir}/{subdirectory}/
         """
         self.log_dir = Path(base_dir) / subdirectory
         self.log_dir.mkdir(parents=True, exist_ok=True)
-    
+
     @staticmethod
     def compute_task_hash(task: Any, length: int = 8) -> str:
         """Compute a hash from the task data.
-        
+
         Args:
             task: The task dictionary or data
             length: Length of the hash to use (default 8 chars)
-        
+
         Returns:
             Hash string
         """
         # Convert task to a stable string representation
         task_str = json.dumps(task, sort_keys=True, default=str)
         # Compute SHA256 hash
-        hash_obj = hashlib.sha256(task_str.encode('utf-8'))
+        hash_obj = hashlib.sha256(task_str.encode("utf-8"))
         # Return first `length` characters of hex digest
         return hash_obj.hexdigest()[:length]
-    
+
     def get_step_dir(self, step: int, mode: str = "train", epoch: int = 0) -> Path:
         """Get the directory path for a specific training or validation step.
-        
+
         Args:
             step: Current training/validation step
             mode: Mode identifier ('train' or 'val'), defaults to 'train'
             epoch: Current epoch number, defaults to 0
-        
+
         Returns:
             Path object for the step directory
         """
         step_dir = self.log_dir / f"{mode}_step_{step}_epoch_{epoch}"
         step_dir.mkdir(parents=True, exist_ok=True)
         return step_dir
-    
+
     def get_episode_filename(self, episode: Episode, step: int) -> str:
         """Generate filename for an episode.
-        
+
         Format: episode_hash{task_hash}_id{episode_id}.json
-        
+
         Args:
             episode: The episode to save
             step: Current training step (not used in filename, but kept for compatibility)
-        
+
         Returns:
             Filename string
         """
         task_hash = self.compute_task_hash(episode.task)
         # Clean episode_id to make it filesystem-safe
-        episode_id_safe = str(episode.id).replace(':', '_').replace('/', '_')
-        
+        episode_id_safe = str(episode.id).replace(":", "_").replace("/", "_")
+
         filename = f"episode_hash{task_hash}_id{episode_id_safe}.json"
         return filename
-    
+
     def log_episode(self, episode: Episode, step: int, mode: str = "train", epoch: int = 0):
         """Log a single episode to its own JSON file in a step-specific directory.
-        
+
         Args:
             episode: The episode to log
             step: Current training/validation step
             mode: Mode identifier ('train' or 'val'), defaults to 'train'
             epoch: Current epoch number, defaults to 0
         """
-        episode_data = {
-            'training_step': step,
-            'epoch': epoch,
-            'episode_id': episode.id,
-            'task': episode.task,
-            'task_hash': self.compute_task_hash(episode.task),
-            'is_correct': episode.is_correct,
-            'termination_reason': episode.termination_reason.value if episode.termination_reason else None,
-            'metrics': episode.metrics,
-            'timing': episode.info.get('timing', {}),
-            'trajectories': []
-        }
-        
+        episode_data = {"training_step": step, "epoch": epoch, "episode_id": episode.id, "task": episode.task, "task_hash": self.compute_task_hash(episode.task), "is_correct": episode.is_correct, "termination_reason": episode.termination_reason.value if episode.termination_reason else None, "metrics": episode.metrics, "timing": episode.info.get("timing", {}), "trajectories": []}
+
         for traj in episode.trajectories:
             traj_data = {
-                'name': traj.name,
-                'uid': traj.uid,
-                'reward': traj.reward,
-                'num_steps': len(traj.steps),
-                'timing': traj.info.get('timing', {}),
-                'steps': [
+                "name": traj.name,
+                "uid": traj.uid,
+                "reward": traj.reward,
+                "num_steps": len(traj.steps),
+                "timing": traj.info.get("timing", {}),
+                "steps": [
                     {
-                        'observation': step.observation,
-                        'thought': step.thought,
-                        'action': step.action,
-                        'reward': step.reward,
-                        'done': step.done,
-                        'model_response': step.model_response,
-                        'chat_completions': step.chat_completions,
-                        'timing': step.info.get('timing', {}),  # Add step-level timing
+                        "observation": step.observation,
+                        "thought": step.thought,
+                        "action": step.action,
+                        "reward": step.reward,
+                        "done": step.done,
+                        "model_response": step.model_response,
+                        "chat_completions": step.chat_completions,
+                        "timing": step.info.get("timing", {}),  # Add step-level timing
                     }
                     for step in traj.steps
-                ]
+                ],
             }
-            episode_data['trajectories'].append(traj_data)
-        
+            episode_data["trajectories"].append(traj_data)
+
         # Write to individual file in step-specific directory
         step_dir = self.get_step_dir(step, mode, epoch)
         filename = self.get_episode_filename(episode, step)
         filepath = step_dir / filename
-        
+
         try:
-            with open(filepath, 'w') as f:
+            with open(filepath, "w") as f:
                 json_str = json.dumps(episode_data, indent=2, default=str)
-                f.write(json_str + '\n')
+                f.write(json_str + "\n")
                 f.flush()  # Ensure data is written to disk
         except Exception as e:
             print(f"Error writing episode to {filepath}: {e}")
             import traceback
+
             traceback.print_exc()
             raise
-    
+
     def log_episodes(self, episodes: list[Episode], step: int, mode: str = "train", epoch: int = 0):
         """Log multiple episodes, each to its own file.
-        
+
         Args:
             episodes: List of episodes to log
             step: Current training/validation step
@@ -148,14 +140,14 @@ def log_episodes(self, episodes: list[Episode], step: int, mode: str = "train",
         for i, episode in enumerate(episodes):
             try:
                 self.log_episode(episode, step, mode, epoch)
-                print(f"[EpisodeLogger] Successfully logged episode {i+1}/{len(episodes)}: {episode.id}")
+                print(f"[EpisodeLogger] Successfully logged episode {i + 1}/{len(episodes)}: {episode.id}")
             except Exception as e:
-                print(f"[EpisodeLogger] Failed to log episode {i+1}/{len(episodes)}: {e}")
+                print(f"[EpisodeLogger] Failed to log episode {i + 1}/{len(episodes)}: {e}")
                 raise
-    
+
     def log_episodes_batch(self, episodes: list[Episode], step: int, mode: str = "train", epoch: int = 0, batch_summary: bool = True):
         """Log multiple episodes and optionally create a batch summary in step-specific directory.
-        
+
         Args:
             episodes: List of episodes to log
             step: Current training/validation step
@@ -165,28 +157,24 @@ def log_episodes_batch(self, episodes: list[Episode], step: int, mode: str = "tr
         """
         # Log individual episodes
         self.log_episodes(episodes, step, mode, epoch)
-        
+
         # Optionally create batch summary in step-specific directory
         if batch_summary and episodes:
             summary_data = {
-                'training_step': step,
-                'epoch': epoch,
-                'mode': mode,
-                'num_episodes': len(episodes),
-                'episode_files': [
-                    self.get_episode_filename(ep, step) for ep in episodes
-                ],
-                'summary_stats': {
-                    'total_correct': sum(1 for ep in episodes if ep.is_correct),
-                    'total_incorrect': sum(1 for ep in episodes if not ep.is_correct),
-                    'accuracy': sum(1 for ep in episodes if ep.is_correct) / len(episodes) if episodes else 0,
-                    'avg_trajectories_per_episode': sum(len(ep.trajectories) for ep in episodes) / len(episodes) if episodes else 0,
-                }
+                "training_step": step,
+                "epoch": epoch,
+                "mode": mode,
+                "num_episodes": len(episodes),
+                "episode_files": [self.get_episode_filename(ep, step) for ep in episodes],
+                "summary_stats": {
+                    "total_correct": sum(1 for ep in episodes if ep.is_correct),
+                    "total_incorrect": sum(1 for ep in episodes if not ep.is_correct),
+                    "accuracy": sum(1 for ep in episodes if ep.is_correct) / len(episodes) if episodes else 0,
+                    "avg_trajectories_per_episode": sum(len(ep.trajectories) for ep in episodes) / len(episodes) if episodes else 0,
+                },
             }
-            
+
             step_dir = self.get_step_dir(step, mode, epoch)
             summary_file = step_dir / "batch_summary.json"
-            with open(summary_file, 'w') as f:
+            with open(summary_file, "w") as f:
                 json.dump(summary_data, f, indent=2)
-
-