pipeline to process Episodes into TrajectoryGroup instead of regrouping

Author: jeffreysijuntan

examples/solver_judge_tinker/train_solver_judge_flow_tinker.sh

@@ -11,7 +11,7 @@ python -m examples.solver_judge_tinker.train_solver_judge_flow_tinker \
     sampling.top_p=1.0 \
     algorithm.adv_estimator=grpo \
     algorithm.norm_adv_by_std_in_grpo=true \
-    algorithm.grouping_level=step \
+    algorithm.grouping_level=trajectory \
     data.max_prompt_length=2048 \
     data.max_response_length=1024 \
     data.train_batch_size=64 \

rllm/trainer/tinker/tinker_agent_trainer.py

@@ -23,7 +23,6 @@
 
 from rllm.agents.agent import Episode, Step, Trajectory
 from rllm.engine.agent_execution_engine import AsyncAgentExecutionEngine
-from rllm.trainer.tinker.tinker_data_processor import TrajectoryGroup
 from rllm.trainer.tinker.tinker_metrics_utils import (
     compute_training_metrics,
     print_episodes,
@@ -215,13 +214,9 @@ async def _fit_agent_async(self):
 
                     logger.info(f"Training for batch {batch_idx}, minibatch {minibatch_count}/{num_minibatches}")
 
-                    # Convert episodes to trajectory groups
-                    # For agent trainer, each episode becomes one group (simple conversion)
-                    trajectory_groups = [TrajectoryGroup(trajectories=episode.trajectories, group_id=episode.id if hasattr(episode, "id") else f"group_{i}") for i, episode in enumerate(minibatch_episodes)]
-
                     # Train immediately (streaming), only optimize on last minibatch
                     t_train_start = time.time()
-                    logprobs, datums = await self.trainer.step(trajectory_groups, learning_rate=learning_rate, beta1=beta1, beta2=beta2, eps=eps, optimizer_step=False)
+                    logprobs, datums = await self.trainer.step(minibatch_episodes, learning_rate=learning_rate, beta1=beta1, beta2=beta2, eps=eps, optimizer_step=False)
                     forward_backward_times.append(time.time() - t_train_start)
                     training_logprobs.extend(logprobs)
                     training_datums.extend(datums)

rllm/trainer/tinker/tinker_data_processor.py

@@ -396,6 +396,96 @@ def to_datum(self) -> tinker.Datum:
         return datums
 
 
+def process_episodes(
+    episodes: list,
+    advantage_computer: TinkerAdvantageComputer,
+    trajectory_filter: TinkerTrajectoryFilter,
+    algorithm_config,
+) -> list[tinker.Datum]:
+    """
+    Main pipeline to convert Episode objects to training datums.
+
+    This function:
+    1. Groups trajectories based on grouping_level configuration
+    2. Computes advantages for each group
+    3. Builds Tinker Datums for training
+
+    Grouping levels:
+    - trajectory: Group trajectories by (task_id, trajectory_name) for multi-agent workflows.
+                 Advantage computed across trajectory rewards.
+    - step: Group individual steps at same position for step-level advantage computation.
+    - episode: Each episode's trajectories form one group (simple single-agent case).
+
+    Args:
+        episodes: List of Episode objects
+        advantage_computer: Computer for calculating advantages
+        trajectory_filter: Filter for removing constant-reward groups
+        algorithm_config: Configuration with grouping_level setting
+
+    Returns:
+        List of Tinker Datum objects ready for training
+    """
+    from collections import defaultdict
+
+    grouping_level = algorithm_config.get("grouping_level", "episode")
+
+    # Group trajectories based on grouping_level
+    trajectory_groups_dict = defaultdict(list)
+
+    def get_task_id(episode):
+        """Extract task_id from episode.id (format: task_id:rollout_idx)"""
+        return ":".join(episode.id.split(":")[:-1]) if ":" in episode.id else episode.id
+
+    if grouping_level == "trajectory":
+        # Group by (task_id, trajectory_name) - for multi-agent workflows like solver-judge
+        for episode in episodes:
+            task_id = get_task_id(episode)
+            for trajectory in episode.trajectories:
+                group_key = (task_id, trajectory.name)
+                trajectory_groups_dict[group_key].append(trajectory)
+
+    elif grouping_level == "step":
+        # Group by (task_id, trajectory_name, step_idx) - for step-level advantages
+        for episode in episodes:
+            task_id = get_task_id(episode)
+            for trajectory in episode.trajectories:
+                for step_idx, step in enumerate(trajectory.steps):
+                    group_key = (task_id, trajectory.name, step_idx)
+                    # Create single-step trajectory
+                    from rllm.agents.agent import Trajectory
+
+                    single_step_traj = Trajectory(steps=[step], reward=step.reward, name=trajectory.name)
+                    trajectory_groups_dict[group_key].append(single_step_traj)
+
+    else:  # "episode" or default
+        # Simple grouping: all trajectories in an episode form one group
+        for episode in episodes:
+            group_key = episode.id
+            trajectory_groups_dict[group_key].extend(episode.trajectories)
+
+    # Convert dict to list of TrajectoryGroup objects for filtering
+    trajectory_groups = [TrajectoryGroup(trajectories=trajs, group_id=str(key)) for key, trajs in trajectory_groups_dict.items()]
+
+    # Apply filtering based on configuration
+    filtered_groups = trajectory_filter.filter_groups(trajectory_groups)
+
+    training_datums = []
+    for group in filtered_groups:
+        # Extract rewards for the group (from all trajectories)
+        group_rewards = [traj.reward for traj in group.trajectories]
+
+        # Compute advantages
+        advantages = advantage_computer.compute(group_rewards)
+
+        # Create datums for all trajectories in the group
+        for trajectory, advantage in zip(group.trajectories, advantages, strict=False):
+            # Use trajectory-level building (merges steps when possible)
+            new_datums = TinkerDatumBuilder.build_datum_from_trajectory(trajectory, advantage)
+            training_datums.extend(new_datums)
+
+    return training_datums
+
+
 def process_trajectory_groups(
     groups: list[TrajectoryGroup],
     advantage_computer: TinkerAdvantageComputer,

rllm/trainer/tinker/tinker_policy_trainer.py

@@ -17,8 +17,7 @@
 from rllm.trainer.tinker.tinker_data_processor import (
     TinkerAdvantageComputer,
     TinkerTrajectoryFilter,
-    TrajectoryGroup,
-    process_trajectory_groups,
+    process_episodes,
 )
 
 if TYPE_CHECKING:
@@ -127,26 +126,23 @@ def _remove_mask(self, datum: tinker.Datum) -> tinker.Datum:
 
     async def step(
         self,
-        groups: list[TrajectoryGroup],
+        episodes: list,
         learning_rate: float = None,
         beta1: float = 0.9,
         beta2: float = 0.95,
         eps: float = 1e-8,
         optimizer_step: bool = True,
     ) -> tuple[list[torch.Tensor], list[tinker.Datum]]:
         """
-        Complete training step: process trajectory groups and update policy.
+        Complete training step: process episodes and update policy.
 
         This method:
-        1. Converts episodes to trajectory groups if needed
-        2. Filters groups (if configured)
-        3. Computes advantages
-        4. Converts to datums
-        5. Performs forward-backward pass
-        6. Applies optimizer step
+        1. Processes episodes to compute advantages and create datums
+        2. Performs forward-backward pass
+        3. Applies optimizer step
 
         Args:
-            data: List of Episode or TrajectoryGroup objects
+            episodes: List of Episode objects
             learning_rate: Learning rate (uses config value if None)
             optimizer_step: Whether to apply optimizer step after forward-backward
 
@@ -159,10 +155,11 @@ async def step(
             learning_rate = self.config.training.learning_rate
 
         # Step 1: Process to datums (includes filtering and advantage computation)
-        training_datums = process_trajectory_groups(
-            groups,
+        training_datums = process_episodes(
+            episodes,
             self.advantage_computer,
             self.trajectory_filter,
+            self.config.algorithm,
         )
 
         # Step 3: Remove mask from datums (not needed by forward_backward)
@@ -199,11 +196,12 @@ async def step(
         # Return both logprobs and datums (with masks for metrics)
         return training_logprobs_D, training_datums
 
-    async def forward_backward_future(self, groups: list[TrajectoryGroup]):
-        training_datums = process_trajectory_groups(
-            groups,
+    async def forward_backward_future(self, episodes: list):
+        training_datums = process_episodes(
+            episodes,
             self.advantage_computer,
             self.trajectory_filter,
+            self.config.algorithm,
         )
 
         datums_no_mask = [self._remove_mask(datum) for datum in training_datums]

rllm/trainer/tinker/tinker_workflow_trainer.py

@@ -16,11 +16,10 @@
 import torch
 from transformers import AutoTokenizer
 
-from rllm.agents.agent import Episode, Trajectory
+from rllm.agents.agent import Episode
 from rllm.engine.agent_workflow_engine import AgentWorkflowEngine
 from rllm.engine.rollout.tinker_engine import TinkerEngine
 from rllm.trainer.tinker.tinker_agent_trainer import TinkerAgentTrainer
-from rllm.trainer.tinker.tinker_data_processor import TrajectoryGroup
 from rllm.trainer.tinker.tinker_policy_trainer import TinkerPolicyTrainer
 
 if TYPE_CHECKING:
@@ -119,15 +118,15 @@ def init_envs_and_agents(self, batch_data):
         self.current_batch = batch_data
 
     async def validate_agent(self, dataloader, sampling_client):
-        all_trajectory_groups = []
+        all_episodes = []
         all_episode_metrics = {}  # episode_id -> episode.metrics dict
         self.agent_execution_engine.rollout_engine.set_sampling_client(sampling_client)
         for batch in dataloader:
             batch = self.build_interleave_batch(batch, 1)
             self.init_envs_and_agents(batch)
-            # For validation, collect all trajectory groups from generator
-            async for trajectory_groups, episode_metrics in self.generate_agent_episodes(group_size=1, minibatch_size=1, return_metrics=True):
-                all_trajectory_groups.extend(trajectory_groups)
+            # For validation, collect all episodes from generator
+            async for episodes, episode_metrics in self.generate_agent_episodes(group_size=1, minibatch_size=1, return_metrics=True):
+                all_episodes.extend(episodes)
                 all_episode_metrics.update(episode_metrics)
 
         # Collect workflow metrics per episode (deduplicated by episode.id)
@@ -138,10 +137,10 @@ async def validate_agent(self, dataloader, sampling_client):
                 for key, value in episode_metric_dict.items():
                     workflow_metrics[key].append(float(value))
 
-        # Compute trajectory-level statistics from all groups
+        # Compute trajectory-level statistics from all episodes
         all_trajectories = []
-        for group in all_trajectory_groups:
-            all_trajectories.extend(group.trajectories)
+        for episode in all_episodes:
+            all_trajectories.extend(episode.trajectories)
 
         mean_reward = sum([traj.reward for traj in all_trajectories]) / len(all_trajectories)
         std_reward = sum([(traj.reward - mean_reward) ** 2 for traj in all_trajectories]) / len(all_trajectories)
@@ -165,17 +164,14 @@ async def validate_agent(self, dataloader, sampling_client):
 
     async def generate_agent_episodes(self, timing_raw=None, meta_info=None, group_size=None, minibatch_size=None, return_metrics=False):
         """
-        Generate trajectory groups in minibatches with overlapping generation and training.
-
-        This uses a background producer task to continuously generate episodes (from rollout)
-        and regroups them into TrajectoryGroup objects for advantage computation.
+        Generate episodes from workflow execution.
 
         Args:
-            return_metrics: If True, yields (trajectory_groups, metrics) tuple where metrics is
-                          {episode_id: {metric_name: value, ...}}. If False, yields only trajectory_groups.
+            return_metrics: If True, yields (episodes, metrics) tuple where metrics is
+                          {episode_id: {metric_name: value, ...}}. If False, yields only episodes.
 
         Yields:
-            list[TrajectoryGroup] or tuple[list[TrajectoryGroup], dict] depending on return_metrics
+            list[Episode] or tuple[list[Episode], dict] depending on return_metrics
         """
 
         num_minibatches = self.config.training.num_minibatches
@@ -187,116 +183,21 @@ async def generate_agent_episodes(self, timing_raw=None, meta_info=None, group_s
 
         episodes = await self.agent_execution_engine.execute_tasks(current_batch, task_ids)
         episodes = self.make_sure_contain_token_and_logprob(episodes)
-        trajectory_groups, episode_metrics = self.regroup(episodes)
-
-        if return_metrics:
-            yield trajectory_groups, episode_metrics
-        else:
-            yield trajectory_groups
-
-    def regroup(self, episodes: list[Episode]) -> tuple[list[TrajectoryGroup], dict]:
-        """
-        Regroup episodes into TrajectoryGroup objects based on grouping_level configuration.
-
-        The grouping level determines how advantages are computed:
-
-        - trajectory: Group trajectories by (task_id, trajectory_name). Each trajectory
-                     keeps all its steps. Trajectories with different names are grouped
-                     separately (important for multi-agent scenarios). Advantage is
-                     computed per trajectory (e.g., via GRPO across trajectory rewards),
-                     then broadcast to all steps in that trajectory during datum creation.
 
-        - step: Group individual steps at the same position (task_id + trajectory_name
-                + step_idx) across different rollouts. Each step becomes a single-step
-                trajectory in a group. Advantage is computed per step (e.g., via GRPO
-                across step rewards).
-
-        The resulting TrajectoryGroup objects are consumed by process_trajectory_groups() which:
-        1. Extracts rewards from trajectories in each group
-        2. Computes advantages across those rewards
-        3. Assigns each trajectory its computed advantage
-        4. Broadcasts the advantage to all steps in the trajectory
+        # Update trajectory-level rewards from step-level rewards
+        for episode in episodes:
+            for trajectory in episode.trajectories:
+                if trajectory.reward == 0.0 and trajectory.steps:
+                    # Compute trajectory reward from step rewards
+                    trajectory.reward = sum(step.reward if step.reward is not None else 0.0 for step in trajectory.steps)
 
-        Args:
-            episodes: List of episodes to regroup
+        # Extract episode metrics if available
+        episode_metrics = {ep.id: ep.metrics for ep in episodes if hasattr(ep, "metrics") and ep.metrics}
 
-        Returns:
-            Tuple of (trajectory_groups, metrics_dict)
-        """
-        grouping_level = self.config.algorithm.grouping_level
-        trajectory_groups = []
-        metrics = {}
-
-        def get_task_id(episode: Episode):
-            return ":".join(episode.id.split(":")[:-1])
-
-        if grouping_level == "trajectory":
-            # Group trajectories by (task_id, trajectory_name)
-            # This ensures trajectories with different names are grouped separately
-            temp_groups = defaultdict(list)
-
-            for episode in episodes:
-                if episode.id not in metrics and episode.metrics:
-                    metrics[episode.id] = episode.metrics
-                task_id = get_task_id(episode)
-
-                # Add all trajectories to the group for this (task_id, trajectory_name)
-                for trajectory in episode.trajectories:
-                    # Each trajectory keeps all its steps
-                    # Compute trajectory-level reward as the sum/mean of step rewards
-                    traj_reward = trajectory.reward if trajectory.reward is not None else sum(step.reward for step in trajectory.steps)
-                    # Update trajectory with proper reward
-                    updated_trajectory = Trajectory(steps=trajectory.steps, reward=traj_reward, name=trajectory.name)
-                    # Group by both task_id and trajectory name
-                    group_key = (task_id, trajectory.name)
-                    temp_groups[group_key].append(updated_trajectory)
-
-            # Create TrajectoryGroup objects from grouped trajectories
-            for group_key, trajectories in temp_groups.items():
-                group_id = f"{group_key[0]}:{group_key[1]}"  # "task_id:trajectory_name"
-                trajectory_group = TrajectoryGroup(trajectories=trajectories, group_id=group_id)
-                trajectory_groups.append(trajectory_group)
-
-            print("Trajectory-level grouping:")
-            print(f"  len episodes: {len(episodes)}")
-            print(f"  len unique (task_id, traj_name) groups: {len(temp_groups)}")
-            print(f"  len trajectory_groups: {len(trajectory_groups)}")
-
-        elif grouping_level == "step":
-            # Group individual steps by step position
-            unique_step_uids = set()
-            unique_task_ids = set()
-            step_groupby_step_uid = defaultdict(list)
-
-            for episode in episodes:
-                if episode.id not in metrics and episode.metrics:
-                    metrics[episode.id] = episode.metrics
-                task_id = get_task_id(episode)
-                unique_task_ids.add(task_id)
-
-                for trajectory in episode.trajectories:
-                    for step_idx, step in enumerate(trajectory.steps):
-                        step_uid = f"{task_id}:{trajectory.name}:{step_idx}"
-                        if step_uid not in unique_step_uids:
-                            unique_step_uids.add(step_uid)
-
-                        step_groupby_step_uid[step_uid].append(step)
-
-            # Create TrajectoryGroup objects where each trajectory contains a single step
-            for step_uid, steps in step_groupby_step_uid.items():
-                trajectories = [Trajectory(steps=[step], reward=step.reward) for step in steps]
-                trajectory_group = TrajectoryGroup(trajectories=trajectories, group_id=step_uid)
-                trajectory_groups.append(trajectory_group)
-
-            print("Step-level grouping:")
-            print(f"  len episodes: {len(episodes)}")
-            print(f"  len unique_task_ids: {len(unique_task_ids)}")
-            print(f"  len unique_step_uids: {len(unique_step_uids)}")
-            print(f"  len trajectory_groups: {len(trajectory_groups)}")
+        if return_metrics:
+            yield episodes, episode_metrics
         else:
-            raise ValueError(f"Invalid grouping_level: {grouping_level}. Must be 'trajectory' or 'step'")
-
-        return trajectory_groups, metrics
+            yield episodes
 
     def make_sure_contain_token_and_logprob(self, episodes: list[Episode]) -> list[Episode]:
         for episode in episodes: