feat: Backend-First Training API (Phase 1) (#521)

* feat: Backend-First Training API (Phase 1)

BREAKING CHANGE: Training is now initiated through backend.train() instead of model.train()

## Migration Guide

### Before (deprecated):
```python
await model.train(trajectory_groups, config=art.TrainConfig(learning_rate=5e-6))
```

### After (recommended):
```python
await model.log(trajectory_groups, split='train')  # Log trajectories
result = await backend.train(model, trajectory_groups, learning_rate=5e-6)
await model.log(metrics=result.metrics, step=result.step, split='train')  # Log training metrics
```

## Key Changes

- **New API**: `backend.train(model, trajectory_groups, **kwargs)` with explicit, type-safe parameters
- **Explicit logging**: `backend.train()` does NOT automatically log trajectories or metrics
- **Extended model.log()**: Now accepts `metrics` and `step` kwargs for logging training metrics directly
- **Structured returns**: `LocalTrainResult` and `ServerlessTrainResult` with step, metrics, and backend-specific fields
- **Fixed get_inference_name()**: Now correctly returns `model.name@step` for LocalBackend
- **Deprecation warning**: `model.train()` emits a warning with migration instructions

## Phase 2 (Future)

In a future release, we will:
- Remove `model.train()` method entirely
- Remove `art.TrainConfig` and `art.dev.TrainConfig` classes

Closes #519

* refactor: simplify examples to single model.log() after training

Combined trajectory and metrics logging into a single call:
  result = await backend.train(model, groups, ...)
  await model.log(groups, metrics=result.metrics, step=result.step, split='train')

Removed redundant comments and pre-training log calls.

---------

Co-authored-by: Cursor Bot <bot@cursor.com>

Author: corbt

examples/2048/train.py

@@ -71,9 +71,9 @@ async def train():
             model,
         )
 
-        await model.train(
-            train_groups,
-            config=art.TrainConfig(learning_rate=1e-5),
+        result = await backend.train(model, train_groups, learning_rate=1e-5)
+        await model.log(
+            train_groups, metrics=result.metrics, step=result.step, split="train"
         )
 
 

examples/benchmarking_comparison_models.py

@@ -127,7 +127,8 @@ async def train_model(model: art.TrainableModel):
             )
             for scenario in batch.items
         )
-        await model.train(groups)
+        result = await backend.train(model, groups)
+        await model.log(groups, metrics=result.metrics, step=result.step, split="train")
 
         if batch.step % 20 == 0:
             # Every 20 steps let's benchmark our model under training so we can

examples/hn_title_generator/train.py

@@ -325,9 +325,11 @@ async def main():
             )
             continue
 
-        await model.train(
-            valid_train_groups,
-            config=art.TrainConfig(learning_rate=LEARNING_RATE),
+        result = await backend.train(
+            model, valid_train_groups, learning_rate=LEARNING_RATE
+        )
+        await model.log(
+            valid_train_groups, metrics=result.metrics, step=result.step, split="train"
         )
 
         if batch.step > 0 and batch.step % EVAL_STEPS == 0:

examples/just-the-facts/just_the_facts/train.py

@@ -81,13 +81,13 @@ async def train(
             ),
         )
 
-        await model.train(
+        result = await backend.train(
+            model,
             groups,
-            config=art.TrainConfig(learning_rate=model.config.learning_rate),
-            _config=art.dev.TrainConfig(
-                scale_rewards=model.config.scale_rewards,
-            ),
+            learning_rate=model.config.learning_rate,
+            scale_rewards=model.config.scale_rewards,
         )
+        await model.log(groups, metrics=result.metrics, step=result.step, split="train")
 
         await backend._experimental_push_to_s3(model)
 

examples/mcp-rl/mcp_rl/train.py

@@ -168,7 +168,8 @@ async def train_mcp_agent(model: art.TrainableModel, use_skypilot: bool = False)
             await model.log(val_groups, split="val")
 
         print("starting train")
-        await model.train(groups, config=art.TrainConfig(learning_rate=learning_rate))
+        result = await backend.train(model, groups, learning_rate=learning_rate)
+        await model.log(groups, metrics=result.metrics, step=result.step, split="train")
 
         await backend._experimental_push_to_s3(
             model,

examples/openenv_echo.py

@@ -86,7 +86,8 @@ async def main() -> None:
             [art.TrajectoryGroup(rollout(model, env_client) for env_client in env_pool)]
         )
 
-        await model.train(groups)
+        result = await backend.train(model, groups)
+        await model.log(groups, metrics=result.metrics, step=result.step, split="train")
 
 
 asyncio.run(main())

examples/prisoners-dilemma.ipynb

@@ -136,17 +136,18 @@
     "    )\n",
     "    await model.log([ts[0] for ts in base_play_trajectories], split=\"versus-base\")\n",
     "    await model.log([ts[1] for ts in base_play_trajectories], split=\"base-model\")\n",
-    "    # Train the model on self-play and base-play trajectories.\n",
-    "    await model.train(\n",
-    "        trajectory_groups=[\n",
-    "            # Since all self-play games have the same starting state and are symmetric, we can gather\n",
-    "            # trajectories from all self-play games into a single trajectory group.\n",
-    "            art.TrajectoryGroup(t for ts in self_play_trajectories for t in ts),\n",
-    "            # We can also gather all base-play _trained model_ trajectories into a single trajectory group.\n",
-    "            # We don't want to train on base model trajectories, because they are sampled from a different distribution.\n",
-    "            art.TrajectoryGroup(ts[0] for ts in base_play_trajectories),\n",
-    "        ],\n",
-    "        config=art.TrainConfig(learning_rate=5e-5),\n",
+    "    # Train the model on self-play and base-play trajectories using the backend-first API.\n",
+    "    trajectory_groups = [\n",
+    "        # Since all self-play games have the same starting state and are symmetric, we can gather\n",
+    "        # trajectories from all self-play games into a single trajectory group.\n",
+    "        art.TrajectoryGroup(t for ts in self_play_trajectories for t in ts),\n",
+    "        # We can also gather all base-play _trained model_ trajectories into a single trajectory group.\n",
+    "        # We don't want to train on base model trajectories, because they are sampled from a different distribution.\n",
+    "        art.TrajectoryGroup(ts[0] for ts in base_play_trajectories),\n",
+    "    ]\n",
+    "    result = await backend.train(model, trajectory_groups, learning_rate=5e-5)\n",
+    "    await model.log(\n",
+    "        trajectory_groups, metrics=result.metrics, step=result.step, split=\"train\"\n",
     "    )"
    ]
   }
@@ -172,4 +173,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

examples/rock-paper-tool-use.ipynb

@@ -52,7 +52,8 @@
     "model = art.TrainableModel(\n",
     "    name=MODEL_NAME, project=\"rock-paper-tool-use\", base_model=BASE_MODEL\n",
     ")\n",
-    "await model.register(LocalBackend())\n",
+    "backend = LocalBackend()\n",
+    "await model.register(backend)\n",
     "client = model.openai_client()\n",
     "\n",
     "\n",
@@ -180,10 +181,10 @@
     "    trajectories = await art.gather_trajectories(\n",
     "        (rollout() for _ in range(64)), max_exceptions=64\n",
     "    )\n",
-    "    await model.train(\n",
-    "        [art.TrajectoryGroup(trajectories)],\n",
-    "        config=art.TrainConfig(learning_rate=5e-5),\n",
-    "    )"
+    "    # Log trajectories and train using the backend-first API\n",
+    "    groups = [art.TrajectoryGroup(trajectories)]\n",
+    "    result = await backend.train(model, groups, learning_rate=5e-5)\n",
+    "    await model.log(groups, metrics=result.metrics, step=result.step, split=\"train\")"
    ]
   }
  ],
@@ -208,4 +209,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

examples/temporal_clue/temporal-clue-7b-async.ipynb

@@ -156,10 +156,11 @@
     "            for trajectory in group:\n",
     "                trajectory.metrics[\"max_reward\"] = max_reward\n",
     "        await model.delete_checkpoints()\n",
-    "        await model.train(\n",
-    "            train_groups,\n",
-    "            config=art.TrainConfig(learning_rate=5e-6),\n",
-    "            _config=art.dev.TrainConfig(precalculate_logprobs=True),\n",
+    "        result = await backend.train(\n",
+    "            model, train_groups, learning_rate=5e-6, precalculate_logprobs=True\n",
+    "        )\n",
+    "        await model.log(\n",
+    "            train_groups, metrics=result.metrics, step=result.step, split=\"train\"\n",
     "        )"
    ]
   }
@@ -185,4 +186,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

examples/temporal_clue/temporal-clue-7b.ipynb

@@ -118,10 +118,15 @@
     "            trajectory.metrics[\"max_reward\"] = max_reward\n",
     "    await model.log(val_groups)\n",
     "    await model.delete_checkpoints()\n",
-    "    await model.train(\n",
+    "    result = await backend.train(\n",
+    "        model,\n",
     "        train_groups,\n",
-    "        config=art.TrainConfig(learning_rate=5e-6),\n",
-    "        _config=art.dev.TrainConfig(precalculate_logprobs=True, scale_rewards=False),\n",
+    "        learning_rate=5e-6,\n",
+    "        precalculate_logprobs=True,\n",
+    "        scale_rewards=False,\n",
+    "    )\n",
+    "    await model.log(\n",
+    "        train_groups, metrics=result.metrics, step=result.step, split=\"train\"\n",
     "    )"
    ]
   }
@@ -147,4 +152,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}