Train agents with Ray and RAGEN

An agent-training experiment combines a model, an environment, a reward definition, and a distributed runtime. Reproducing it requires preserving all four, not just the model weights.

RAGEN provides multi-turn agent reinforcement learning using verl and Ray. Polyaxon creates and tracks the Kubernetes workload; Ray schedules the training and rollout actors. This article uses RAGEN's Sokoban configuration to show how the framework settings relate to a Ray cluster's resources.

Prepare a versioned runtime

Install KubeRay and enable the RayCluster integration on the Polyaxon-managed compute cluster. Before starting a multi-node run, check GPU communication with NCCL.

Build an image from a reviewed RAGEN commit, including its pinned verl submodule, environment dependencies, and compatible inference backend. Follow the repository's installation instructions. Put the installed checkout at /workspace/RAGEN for the component below, or change the working directory to match your image.

Record the Ray version installed in that image. Use it for rayVersion as well as the head and worker containers. Do not select a newer Ray version independently of RAGEN's dependency set.

Configure a ragen-storage connection that mounts shared writable checkpoint storage at /mnt/ragen on the head and every worker. Make model downloads and environment dependencies available to all replicas. Add secret connections if you enable external tracking or use a private model.

Match the training configuration to the cluster

The Sokoban recipe inherits the project's base configuration. Save this component as ragen-ray.yaml; it requests a CPU head and two one-GPU workers, with matching trainer.nnodes and trainer.n_gpus_per_node settings.

version: 1.1
kind: component
name: ragen-sokoban
inputs:
- name: image
  type: str
- name: ray_version
  type: str
plugins:
  shm: true
run:
  kind: raycluster
  connections: [ragen-storage]
  rayVersion: "{{ ray_version }}"
  entrypoint: >-
    python train.py --config-name _2_sokoban
    +ray_kwargs.ray_init.address=auto
    trainer.nnodes=2
    trainer.n_gpus_per_node=1
    trainer.logger=[console]
    trainer.default_local_dir=/mnt/ragen/{{ globals.uuid }}/checkpoints
    trainer.local_log_dir=/mnt/ragen/{{ globals.uuid }}/logs
  head:
    rayStartParams:
      dashboard-host: "0.0.0.0"
      num-gpus: "0"
    container:
      image: "{{ image }}"
      workingDir: /workspace/RAGEN
  workers:
    training-workers:
      replicas: 2
      minReplicas: 2
      maxReplicas: 2
      container:
        image: "{{ image }}"
        workingDir: /workspace/RAGEN
        resources:
          limits:
            nvidia.com/gpu: "1"

This is a deployment configuration for your prepared image, not a bundled RAGEN installation. Set CPU and memory requests through the component or a preset, and select GPUs with enough memory for the model, rollout engine, and trainer. Use placement rules when the experiment specifically needs workers on separate physical nodes.

RAGEN's training entrypoint accepts Ray initialization settings and constructs its resource pool from the trainer's node and GPU counts. The address=auto override connects it to the existing cluster. Check those fields against the commit in your image when upgrading.

Submit the operation

Set RAGEN_IMAGE to your image's immutable registry digest and RAGEN_RAY_VERSION to its installed Ray version:

: "${RAGEN_IMAGE:?Set the digest of your prepared RAGEN image}"
: "${RAGEN_RAY_VERSION:?Set the Ray version installed in that image}"
polyaxon run -f ragen-ray.yaml \
  -P image="$RAGEN_IMAGE" \
  -P ray_version="$RAGEN_RAY_VERSION"
polyaxon ops dashboard

Confirm that both workers join and that the training actors use the intended GPU allocation. More worker pods do not automatically make the framework use more resources: update the trainer configuration and memory assumptions together.

Preserve the experiment, not just its checkpoint

Retain the repository and submodule commits, image digest, model revision, environment configuration, reward settings, seeds, and resolved training configuration. Keep validation episodes separate from training episodes and compare task success alongside training loss and reward.

The example uses console logging, so Polyaxon can capture the framework's output. Console metrics are not automatically structured Polyaxon metrics; use a supported tracking integration or explicit metric logging if you need charts and run comparisons.

Checkpoint paths include the run UUID to avoid collisions. To resume another operation, explicitly point the framework at the retained checkpoint and verify its resume settings. The verl trainer configuration documents these fields for the submodule revision reviewed here.

Investigate failures at the responsible layer

If workers do not join, inspect image access, Ray version alignment, head reachability, and network policy. If actors wait despite available GPUs, compare the framework's requested resource pool with Ray's registered resources. If training runs out of memory, inspect the model, rollout allocation, sequence lengths, and batch sizes before adding workers.

Missing checkpoints are often a storage-path or permission issue. Verify the mount on every replica and confirm that the framework actually saved a checkpoint before the run ended.

Stop the cluster when the experiment is complete and outputs are preserved:

polyaxon ops stop