[RFC] verl with Agent/Env

[eric-haibin-lin]

Posting on behalf of @zw0610

Contemporary Train Pattern

In verl.trainer.ppo.ray_trainer.RayTrainer, the fit(self) method follows this pattern:

1. Fetch a batch (from a parquet file): ray_trainer.py#L826
2. Generate responses as actions: ray_trainer.py#L849
3. Compute advantages
4. Critic loss + critic update
5. Actor loss + actor update
6. Update rollout
   This train pattern is not compatible with Agent/Env RL training, which will be illustrated below.

Agent/Env

There might be thousands of definitions of what Agent is. In this case, we can presume an agent is a proxy with certain capability with help from a LLM inference server to interact with some environment (Game, Github Code Issue, etc.) for some tasks.
The introduction of Agent turns the prompt generation in RLHF closer to traditional RL training and expands the cases that LLM can work with.

Train with Agent/Env

There are two main conflicts between training with Agent and the contemporary train pattern:

1. Separation of train process (the driver process in our single-controller paradigm) and prompt data
   Unlike fetching directly from offline data files, prompt data now is generated in Agents, which usually run as different processes.
   To address this issue, we introduced the class MockedServer to collect prompts and worked as a special Dataset for minimal change to the existing train pattern.
2. Dynamic prompt/train batch generation
   If the Agents/Envs require multiple rounds to get the final score, which means traditional Dataset with determined content and length cannot be used if we wish to overlap generation and train stages. (The long-tail issue becomes more severe with multi-round agents/envs.)
   To address this issue, we introduced the class BufferDataset to serve as a pipe between generation and train stages and still maintain a single-controller paradigm in driver script.

Function Call v.s. Mocked Server

There is discussion on whether making the Agent as Sever waiting for function call from training process or attaching a mocked LLM inference server to collect prompt posted by Agents.
We decide to adopt the mocked server solution given:

1. The high diversity of agents now and in the future, release engineers from endless work to transform the agent to a server every time a new agent emerges;
2. The consistency and code reuse between generation stage and evaluation stage.

Design

Architecture

Flow

Here we explained an async flow. A synchronized training flow can be achieved by a barrier between fetch tasks from parquet file and the update rollout step.
We assume the tasks for agents are stored in a parquet file with a similar format to regular data files.

1. Fetch a batch of tasks (from a parquet file)
2. Enqueue the tasks into ServerWithTask
3. May set this in a thread
4. Sample a batch of prompts from BufferDataset of prompt
5. Generate responses as actions
6. Throw the unioned batch (prompt + response) into MultiRoundBufferDataset
7. Sample a batch of prompt-response-score from MultiRoundBufferDataset
8. Skip this sample action and the following steps if the length of MultiRoundBufferDataset is shorter than the batch size
9. Compute advantages
10. Critic loss + critic update
11. Actor loss + actor update
12. Update rollout

Demo

train_dataset = MultiRoundBufferDataset()
train_dataloader_iter = iter(Dataloader(train_dataset))

mocked_server, server_thread = create_and_launch_task_pool(
                                   notify_score_fn=train_dataset.notify_score)
                                   
prompt_dataset = BufferDataset(buffer=mocked_server.prompt_queue)
prompt_dataloader_iter = iter(Dataloader(prompt_dataset))

# enqueue task into task server
threading.Thread(target=enqueue_task, args=("xxx.parquet",))

# sample prompt batch and generate sequences
batch = next(prompt_dataloader_iter)
responses = rollout.generate_sequences(prompt_batch)
batch.union(responses)
train_dataset.put_batch(batch)

# sample
if train_dataset.buffer.qsize() >= batch_size:
    train_batch = next(prompt_dataloader_iter)
    # do ppo

ongoing PR #808

[yufansong]

I am curious why Agent will fetch task, rather than ServerWithTask push task to agents. Is there any particular reasoning behind our design?

[zw0610]

I am curious why Agent will fetch task, rather than ServerWithTask push task to agents. Is there any particular reasoning behind our design?

@yufansong To save the effort of building a service discovery system for multiple Agent instances.

[weiwangorg]

I am very interested in this feature. I'm a bit unclear about the details of steps 2 and 3 in Flow, specifically how env interacts with multiple rollout workers. Assuming our environment is a card game with three players, each of whom is an LLM based agent. What I understand is that the game environment and control logic are stored in ServerWithTask. But how to start three rollout workers and how does ServerWithTask sequentially notify the three workers to take actions while controlling the game process.

[zw0610]

@weiwangorg
ServerWithTask only provides information prepared to launch a task. Say this is a card game with 3 players (agents), all of which in my understanding should use different rollout workers. Following this design, all requests (prompts) from these 3 agents will be sent to the MockedServer and stored in BufferDataset for Prompt with different model_name field on each request. After a batch of sampled from the BufferDataset for Prompt, the driver process should split the batch into multiple DataProto, each of which shares the same model_name and dispatch theese DataProto to the corresponding rollout workers.

[lei-lei-shanda]

From reading the comment (as it is not clear from the RFC), this feature want to support:

1. an agent's trajectory (action_i, observation_i, reward_i) in the environment can have multiple steps. Right now it is single step i=1 and NO observation_i.
2. training MULTIPLE agents simultaneously. Right now it is single actor.

I wonder what does agent-environment interaction, aka observation_i happens in this arch?

[sanxing-chen]

two quick questions

1. In traditional RL, it’s common to incorporate future rewards in the advantage (e.g., GAE) calculation. However, the proposed method appears to consider only the immediate reward.
2. Since rollouts are generated as one-step transitions in a batch, the overall speed is constrained by the slowest example in the batch. It would be more efficient if multi-step trajectory rollouts could be generated in parallel.

[zw0610]

@sanxing-chen

For the first question, what I think is the opposite. The MultiRoundBufferDataset will put a multi-round trajectory into the buffer only when the episode ends (as the future/final reward comes). I have to admit that in this design, there is no mechanism to receive and hold the immediate reward.

For the second question, we spot this long-tail problem as well. This design does not prohibit multiple steps from being generated in parallel as long as the agent can send multiple prompts concurrently by itself. (In our cases, however, the next round would differ according to the response from the previous round, making it sequentially by nature.