ADu2021/aworld-distributed-agent-training

AWorld: Orchestrating Training Recipe for Agentic AI

Core Concept

AWorld addresses the bottleneck in reinforcement learning for agents: the slow collection of environmental interactions. By distributing task execution across multiple cluster nodes, rather than sequential single-machine execution, the system dramatically accelerates experience generation. This enables practical training of capable agents on complex benchmarks like GAIA, where computational efficiency is critical.

Architecture Overview

• Distributed Executor: Parallel task execution across cluster nodes
• Centralized Experience Aggregation: Collects trajectories from all workers
• Efficient Communication: Minimal overhead for distributed coordination
• Scalable RL Pipeline: From experience collection through model training
• Production-Ready Implementation: Open-source reference implementation

Implementation Steps

Stage 1: Design Distributed Task Execution Framework

Create infrastructure for parallel task distribution and result collection.

# Distributed task execution framework
import asyncio
from typing import Dict, List, Any
import pickle
import queue
from concurrent.futures import ProcessPoolExecutor

class DistributedExecutor:
    """Execute tasks across multiple nodes in parallel"""

    def __init__(self, num_workers: int, max_queue_size: int = 1000):
        self.num_workers = num_workers
        self.executor = ProcessPoolExecutor(max_workers=num_workers)
        self.task_queue = queue.Queue(maxsize=max_queue_size)
        self.result_queue = queue.Queue()
        self.running_tasks = {}

    def submit_task(self, task_id: str, task_fn, *args, **kwargs):
        """Submit a task for distributed execution"""
        future = self.executor.submit(task_fn, *args, **kwargs)
        self.running_tasks[task_id] = future
        return task_id

    def collect_results(self, timeout: float = None) -> Dict:
        """Collect completed task results"""
        completed = {}
        for task_id, future in list(self.running_tasks.items()):
            if future.done():
                try:
                    result = future.result(timeout=timeout)
                    completed[task_id] = {"success": True, "result": result}
                except Exception as e:
                    completed[task_id] = {"success": False, "error": str(e)}
                del self.running_tasks[task_id]

        return completed

    def shutdown(self):
        """Clean shutdown"""
        self.executor.shutdown(wait=True)


class RemoteNodeExecutor:
    """Execute on remote cluster nodes (for distributed setup)"""

    def __init__(self, nodes: List[str]):
        """
        Args:
            nodes: List of remote node addresses (e.g., "worker1.cluster:5000")
        """
        self.nodes = nodes
        self.node_load = {node: 0 for node in nodes}

    async def submit_to_least_loaded(self, task_fn, *args, **kwargs) -> str:
        """Route task to least-loaded node"""
        node = min(self.nodes, key=lambda n: self.node_load[n])
        self.node_load[node] += 1

        task_id = await self.submit_remote_task(node, task_fn, *args, **kwargs)
        return task_id

    async def submit_remote_task(self, node: str, task_fn, *args, **kwargs) -> str:
        """Send task to remote node via RPC"""
        import rpc_client  # Hypothetical RPC library

        client = rpc_client.connect(node)
        serialized_fn = pickle.dumps(task_fn)
        serialized_args = pickle.dumps((args, kwargs))

        task_id = await client.submit_task(
            function=serialized_fn,
            arguments=serialized_args
        )
        return task_id

Stage 2: Implement Task and Trajectory Definition

Define what constitutes a task and how to collect trajectories from execution.

# Task and trajectory definitions
from dataclasses import dataclass
from typing import Optional

@dataclass
class Task:
    """Single task for agent to solve"""
    task_id: str
    instruction: str
    tools: List[str]  # Available tools
    ground_truth: Any  # Expected solution
    metadata: Dict = None

    def to_prompt(self) -> str:
        """Convert task to agent prompt"""
        return f"Instruction: {self.instruction}\nAvailable tools: {self.tools}"


@dataclass
class Trajectory:
    """Single execution trace from task"""
    task_id: str
    agent_id: str
    steps: List[Dict]  # Each step: action, observation, reward
    final_result: Any
    success: bool
    token_count: int
    execution_time: float
    metadata: Dict = None

    def to_dict(self):
        return {
            "task_id": self.task_id,
            "steps": self.steps,
            "success": self.success,
            "token_count": self.token_count,
            "execution_time": self.execution_time
        }

Stage 3: Create Parallel Task Executor with Agent Integration

Execute tasks on distributed workers with agent inference.

# Parallel task execution with agent
import time
from typing import Callable

class ParallelTaskExecutor:
    """Execute tasks in parallel across workers"""

    def __init__(
        self,
        agent_model: Any,
        distributed_executor: DistributedExecutor,
        task_loader: Callable
    ):
        self.agent = agent_model
        self.executor = distributed_executor
        self.task_loader = task_loader
        self.trajectories = []

    def execute_single_task(
        self,
        task: Task,
        max_steps: int = 10,
        timeout: float = 60
    ) -> Trajectory:
        """
        Execute a single task with agent.
        Can run in parallel across workers.
        """
        steps = []
        current_state = {"instruction": task.instruction}
        total_tokens = 0
        start_time = time.time()
        success = False

        try:
            for step_idx in range(max_steps):
                # Agent generates action
                action, action_tokens = self.agent.generate_action(
                    current_state,
                    max_tokens=256
                )
                total_tokens += action_tokens

                # Execute action
                observation = self.execute_action(action, task)

                # Record step
                steps.append({
                    "step": step_idx,
                    "action": action,
                    "observation": observation,
                    "tokens": action_tokens
                })

                # Update state
                current_state["last_action"] = action
                current_state["last_observation"] = observation

                # Check success
                if observation.get("done"):
                    success = observation.get("success", False)
                    break

        except Exception as e:
            steps.append({
                "step": "error",
                "error": str(e)
            })

        execution_time = time.time() - start_time

        trajectory = Trajectory(
            task_id=task.task_id,
            agent_id=self.agent.model_id,
            steps=steps,
            final_result=current_state.get("last_observation"),
            success=success,
            token_count=total_tokens,
            execution_time=execution_time
        )

        return trajectory

    def execute_action(self, action: str, task: Task) -> Dict:
        """Execute an action in the environment"""
        # Parse action (e.g., "use_tool[search] with query[...]")
        tool_name = self.parse_tool_name(action)
        tool_args = self.parse_tool_args(action)

        try:
            if tool_name in task.tools:
                # Simulate or call actual tool
                result = self.call_tool(tool_name, tool_args)
                return {"observation": result, "done": False}
            else:
                return {"observation": f"Unknown tool: {tool_name}", "done": True}

        except Exception as e:
            return {"observation": f"Error: {str(e)}", "done": True}

    def distribute_tasks(self, tasks: List[Task]) -> List[Trajectory]:
        """
        Distribute all tasks across workers.
        This is where the 14.6x speedup comes from!
        """
        # Submit all tasks to distributed executor
        submitted_tasks = {}
        for task in tasks:
            task_id = self.executor.submit_task(
                task.task_id,
                self.execute_single_task,
                task
            )
            submitted_tasks[task_id] = task

        # Collect results as they complete
        trajectories = []
        while submitted_tasks:
            results = self.executor.collect_results(timeout=1.0)

            for task_id, result in results.items():
                if result["success"]:
                    trajectories.append(result["result"])
                else:
                    # Handle failed execution
                    print(f"Task {task_id} failed: {result['error']}")

                del submitted_tasks[task_id]

        return trajectories

Stage 4: Implement Experience Aggregation and RL Training

Collect experiences from all workers and train the model.

# Experience aggregation and RL training
import numpy as np

class ExperienceBuffer:
    """Buffer for collecting trajectories"""

    def __init__(self, max_size: int = 100000):
        self.trajectories = []
        self.max_size = max_size

    def add_trajectory(self, traj: Trajectory):
        """Add trajectory to buffer"""
        self.trajectories.append(traj)
        if len(self.trajectories) > self.max_size:
            self.trajectories.pop(0)

    def add_batch(self, trajs: List[Trajectory]):
        """Add multiple trajectories"""
        for traj in trajs:
            self.add_trajectory(traj)

    def sample_batch(self, batch_size: int) -> List[Trajectory]:
        """Sample random batch"""
        indices = np.random.choice(len(self.trajectories), batch_size)
        return [self.trajectories[i] for i in indices]

    def compute_returns(self, discount: float = 0.99):
        """Compute cumulative returns for trajectories"""
        for traj in self.trajectories:
            returns = []
            cumulative = 0
            for step in reversed(traj.steps):
                reward = step.get("reward", 1.0 if traj.success else 0.0)
                cumulative = reward + discount * cumulative
                returns.insert(0, cumulative)
            traj.returns = returns


class RLTrainer:
    """Train agent using RL on collected trajectories"""

    def __init__(self, agent_model, learning_rate: float = 1e-5):
        self.model = agent_model
        self.optimizer = agent_model.optimizer_class(
            agent_model.parameters(), lr=learning_rate
        )

    def train_step(self, batch: List[Trajectory]) -> float:
        """Single RL training step"""
        total_loss = 0

        for traj in batch:
            # Compute policy gradient loss
            for step_idx, step in enumerate(traj.steps):
                action_tokens = step.get("tokens", 0)
                # In practice, compute log probs of actions taken
                log_prob = self.model.get_log_prob(
                    step["action"],
                    traj.steps[:step_idx]
                )

                # Return as advantage signal
                advantage = traj.returns[step_idx]

                # Policy gradient: maximize log_prob * advantage
                loss = -(log_prob * advantage)
                total_loss += loss

        # Backward pass
        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()

        return (total_loss / len(batch)).item()


class DistributedTrainingPipeline:
    """Full pipeline: distribute execution, collect experiences, train"""

    def __init__(
        self,
        agent_model,
        num_workers: int,
        task_dataset: List[Task]
    ):
        self.agent = agent_model
        self.executor = DistributedExecutor(num_workers)
        self.task_executor = ParallelTaskExecutor(
            agent_model,
            self.executor,
            lambda: task_dataset
        )
        self.experience_buffer = ExperienceBuffer()
        self.trainer = RLTrainer(agent_model)
        self.task_dataset = task_dataset

    def train_epoch(self, tasks_per_epoch: int = 1000):
        """Single training epoch"""
        # Select random tasks
        selected_tasks = np.random.choice(
            self.task_dataset,
            size=min(tasks_per_epoch, len(self.task_dataset)),
            replace=False
        ).tolist()

        print(f"Executing {len(selected_tasks)} tasks in parallel...")

        # Execute all tasks in parallel across workers
        trajectories = self.task_executor.distribute_tasks(selected_tasks)

        print(f"Collected {len(trajectories)} trajectories")

        # Add to experience buffer
        self.experience_buffer.add_batch(trajectories)

        # Compute returns
        self.experience_buffer.compute_returns()

        # Train on batches
        print("Training on collected experiences...")
        batch_size = 32
        for _ in range(len(trajectories) // batch_size):
            batch = self.experience_buffer.sample_batch(batch_size)
            loss = self.trainer.train_step(batch)

        return {
            "trajectories": len(trajectories),
            "success_rate": sum(t.success for t in trajectories) / len(trajectories),
            "avg_tokens": np.mean([t.token_count for t in trajectories])
        }

Practical Guidance

Scaling Configuration

• Worker Count: Start with 8-16 workers; scale up to 64+ for large benchmarks
• Task Batch Size: 100-1000 tasks per distribution cycle
• Experience Buffer: Keep last 10,000-100,000 trajectories for training stability
• Training Frequency: After each 1000 task executions, perform 1-5 training epochs

Performance Optimization

• Network Efficiency: Compress trajectories before transfer (usually 10-100KB each)
• Load Balancing: Use least-loaded routing to avoid hotspots
• Fault Tolerance: Implement retry logic for failed tasks
• Memory Management: Stream results to disk if buffer exceeds RAM

Baseline Metrics

• Single-Node Baseline: ~10 tasks/minute on 1 GPU
• 16-Worker Cluster: ~140-150 tasks/minute (14.6x speedup)
• GAIA Performance: 32.23% pass@1 with Qwen3-32B agent (vs 27.91% GPT-4o)

When to Use

• Training agents on complex benchmarks (GAIA, ScienceBoard)
• Settings with access to multi-node clusters
• Problems requiring millions of task executions
• Scenarios where wall-clock training time matters

When NOT to Use

• Single-machine environments without cluster access
• Tasks with very fast execution (<1 second)
• Strongly sequential dependencies between tasks
• Scenarios with limited network bandwidth

Reference

AWorld: Orchestrating Training Recipe for Agentic AI. arXiv:2508.20404

• https://arxiv.org/abs/2508.20404