feat(slime): experimental slime backend on the rollout gateway by Linbo-Liu · Pull Request #96 · awslabs/agentcore-rl-toolkit

Linbo-Liu · 2026-07-31T22:33:24Z

Issue #, if available:

Description of changes:

Summary

Adds backends/experimental/slime/ — a slime training backend built on the in-repo rollout_gateway. It plugs into stock slime/train.py through slime's public extension points only (--rollout-function-path, --custom-reward-post-process-path, --custom-config-path), so there is no forked trainer and no custom entrypoint. This is the successor to backends/slime, which stays untouched (it still depends on the external rllm-model-gateway) until this graduates.

What's in it

integration/rollout.py — the custom rollout function (generate_rollout).
Serves one RolloutGateway in-process on a background thread (ThreadedGatewayServer), sampling token-in/token-out via SGLang's native /generate (SglangHttpBackend against slime's router) and rendering with the HF chat template of --hf-checkpoint. The gateway owns tokenization in both directions, so loss masking is well-defined and there's no retokenization drift between turns.

integration/rewards.py — normalize_episode_rewards, replacing slime's default reshape-based normalization, which assumes a fixed turn count per row. Aggregates rows to one reward per episode by (group_index, gateway_session_id), normalizes across episodes within a task group, then writes the result back to every row — so a 3-turn success and a 2-turn failure each count as one data point. Strategies are pluggable (grpo default, identity); padding rows (group_index=-1) are skipped.

integration/sglang_parsing.py — derender seams built from SGLang's own detectors. Because the gateway samples via /generate, it bypasses the server's OpenAI layer, so reasoning/tool-call parsing has to happen gateway-side. build_tool_parser (FunctionCallParser) and build_reasoning_parser (ReasoningParser) are wired automatically from slime's existing --sglang-tool-call-parser / --sglang-reasoning-parser args, validated eagerly at gateway construction so a bad parser name fails at startup rather than on the first turn. Each stage is independent — unset one and the gateway keeps its dependency-free default. This lives in the slime backend, not rollout_gateway, because slime trainers always have sglang importable while the gateway package stays engine-free.

Example-agent changes

Every rl_app.py now reads api_key = payload["_rollout"].get("api_key", "EMPTY") and passes it to its model client, so the gateway can key capture off the Bearer slot; "EMPTY" (the vLLM convention) keeps plain evaluation endpoints working unchanged. Also fixes response-text extraction in the officebench/appworld agents: response.message["content"][0]["text"] raises KeyError when the first content block is a tool-use or reasoning block, so text blocks are now filtered and joined.

.gitignore: broadened the local-config ignore to backends/**/examples/**/config.yaml (was slime-specific) and added *.jsonl for trace logs.

Test plan

Validated end to end against a live ACR deployment of examples/strands_math_agent — GSM8K GRPO with Qwen3-0.6B on 8×B200, colocated train+rollout, wandb-logged. Exact versions pinned in SETUP.md → Tested Versions.

Test accuracy increased from 6% -> ~80% after 100 GRPO steps.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

luyuzhe111 · 2026-07-31T23:23:22Z

@@ -0,0 +1,23 @@
+# Wandb configuration — copy to .wandb.env (gitignored) and fill in values.
+#


users should be able to run wandb cli login once and they won't need this API setup. can probably fold a note about wandb in the training script?

luyuzhe111 · 2026-07-31T23:26:11Z

+set -euo pipefail
+
+# === Paths (set these via env) ===
+SLIME_DIR="${SLIME_DIR:?Set SLIME_DIR (path to the slime repo, e.g. /root/slime)}"


hmm do we still need this

This is still needed. pip install slime will not install slime/train.py because it's outside its src dir: https://github.com/THUDM/slime/blob/main/train.py

luyuzhe111 · 2026-07-31T23:31:38Z

+CONFIG="${CONFIG:-$(dirname "$0")/config.yaml}"
+
+# Set your cuda path. CUDA 13 only — cu12 is not supported.
+CUDA_HOME="${CUDA_HOME:-/usr/local/cuda-13.0}"


this should only impact compiler choice during installation? do we actually need this here?

luyuzhe111 · 2026-07-31T23:32:10Z

+  exit 1
+fi
+
+# slime/ray/actor_group.py hardcodes the cu12-named torch_memory_saver preload .so,


wondering if this is still true with the latest slime version

Yes. https://github.com/THUDM/slime/blob/main/slime/ray/actor_group.py#L64-L84

luyuzhe111 · 2026-07-31T23:33:04Z

+source ${SLIME_DIR}/scripts/models/${MODEL_TYPE}.sh
+
+# === Launch training ===
+export no_proxy=127.0.0.1


wondering what does this do

luyuzhe111 · 2026-07-31T23:36:23Z

+from collections import defaultdict
+
+
+def _grpo_normalize(episode_rewards: list[float], std_normalize: bool) -> list[float]:


i'm a bit surprised we actually need to implement this? slime should support black-box agent with reward coming from the session already? can we fold the reward extraction in the rollout function so that we can just reuse reward processing logics from slime?

luyuzhe111 · 2026-07-31T23:38:16Z

+# ---------------------------------------------------------------------------
+
+
+def _import_slime_types():


this try except doesn't make much sense to me.

luyuzhe111 · 2026-07-31T23:39:18Z

+    independent: an unset one leaves that derender stage on the gateway's
+    dependency-free default.
+    """
+    try:


try except also meaningless here?

luyuzhe111 · 2026-07-31T23:41:03Z

+# ---------------------------------------------------------------------------
+
+
+def _make_noop_sample(group_index: int = -1, session_id: str = "", status_name: str = "COMPLETED"):


do we still need this

luyuzhe111 · 2026-07-31T23:44:08Z

+# ephemeral NVMe volume (present on AWS DLAMI/HyperPod nodes) unless the caller
+# already set TMPDIR.
+if [ -z "${TMPDIR:-}" ] && [ -d /opt/dlami/nvme ]; then
+  TMPDIR="/opt/dlami/nvme/${USER}/tmp"


TMPDIR remapping should be in .bashrc

luyuzhe111 · 2026-07-31T23:48:05Z

+uv pip install --reinstall-package pyjwt PyJWT
+
+# Install slime
+git clone https://github.com/THUDM/slime.git


we actually need a local slime copy? :))))

luyuzhe111 · 2026-07-31T23:57:22Z

+git -C Megatron-LM apply --3way "$SLIME_PATCH_DIR/megatron.patch"
+
+# sglang patches: apply into site-packages without leaving cwd (patch -d).
+patch -d "$SITE_PACKAGES" -p2 -F0 -N < "$SLIME_PATCH_DIR/sglang.patch"


what patch do we need exactly?

feat(slime): Add integration for slime + rollout gateway.

ac1a8ba

Linbo-Liu force-pushed the slime-integration branch from baeee13 to ac1a8ba Compare July 31, 2026 22:38

luyuzhe111 reviewed Jul 31, 2026

View reviewed changes

		@@ -0,0 +1,23 @@
		# Wandb configuration — copy to .wandb.env (gitignored) and fill in values.
		#

		from collections import defaultdict


		def _grpo_normalize(episode_rewards: list[float], std_normalize: bool) -> list[float]:

		# ---------------------------------------------------------------------------


		def _import_slime_types():

		# ---------------------------------------------------------------------------


		def _make_noop_sample(group_index: int = -1, session_id: str = "", status_name: str = "COMPLETED"):

Uh oh!

Conversation

Linbo-Liu commented Jul 31, 2026

Summary

What's in it

Example-agent changes

Test plan

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants