diff --git a/.github/workflows/dictionary/reinforcement-learning.txt b/.github/workflows/dictionary/reinforcement-learning.txt
index 36e21cba9..8133bd4fc 100644
--- a/.github/workflows/dictionary/reinforcement-learning.txt
+++ b/.github/workflows/dictionary/reinforcement-learning.txt
@@ -13,6 +13,7 @@ multiproc
ndarray
ocdbt
orbax
+pathwaysjob
prefuse
pyconfig
relpath
diff --git a/README.md b/README.md
index e6fb559e0..a09fa9fcb 100644
--- a/README.md
+++ b/README.md
@@ -70,6 +70,7 @@ the primary runtime.
- [Reinforcement Learning reference architecture](/docs/platforms/gke/base/use-cases/reinforcement-learning/README.md)
- [Single-host reinforcement learning with TPUs using GRPO algorithm](/docs/platforms/gke/base/use-cases/reinforcement-learning/single-host-tpu-grpo/README.md)
+ - [Multi-host reinforcement learning with TPUs using GRPO algorithm](/docs/platforms/gke/base/use-cases/reinforcement-learning/multi-host-tpu-grpo/README.md)
### Guides
diff --git a/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/Dockerfile b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/Dockerfile
new file mode 100644
index 000000000..887f726c8
--- /dev/null
+++ b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/Dockerfile
@@ -0,0 +1,35 @@
+# syntax=docker.io/docker/dockerfile:1.17.1
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+FROM astral/uv:python3.12-bookworm-slim
+
+# Use copy mode instead of hardlinks across filesystems
+ENV UV_LINK_MODE=copy
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+ git \
+ && rm -rf /var/lib/apt/lists/*
+
+# Leverage uv's built-in caching for fast subsequent builds
+RUN --mount=type=cache,target=/root/.cache/uv \
+ uv pip install --system maxtext[tpu-post-train]==0.2.2 --resolution=lowest
+
+# Script name for the extra TPU post-training dependencies
+RUN install_tpu_post_train_extra_deps
+
+COPY --from=primary train.py .
+
+CMD ["python3", "train.py"]
diff --git a/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/cloudbuild.yaml b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/cloudbuild.yaml
new file mode 100644
index 000000000..77ea0257b
--- /dev/null
+++ b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/cloudbuild.yaml
@@ -0,0 +1,31 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+images:
+ - ${_DESTINATION}
+
+options:
+ logging: CLOUD_LOGGING_ONLY
+ machineType: E2_HIGHCPU_8
+
+steps:
+ - args:
+ - build
+ - --build-context=primary=container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src
+ - --file=container-images/tpu/rl-tpu-maxtext-grpo-multi-host/Dockerfile
+ - --tag=${_DESTINATION}
+ - .
+ id: "Build Reinforcement Learning on TPU image"
+ name: "docker.io/docker:28.3.3-dind-alpine3.22"
+ waitFor: ["-"]
diff --git a/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src/train.py b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src/train.py
new file mode 100644
index 000000000..a774ff7ee
--- /dev/null
+++ b/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src/train.py
@@ -0,0 +1,272 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import logging
+import os
+import subprocess
+import sys
+
+import clu.metric_writers
+import jax
+import jax.numpy as jnp
+import mlflow
+from huggingface_hub import login
+from mlflow.tracking import MlflowClient
+
+# Mute the noisy vLLM TPU runner warnings
+logging.getLogger("tpu_runner").setLevel(logging.ERROR)
+
+# --- 1. SYSTEM & CACHING ---
+os.environ.update(
+ {
+ "PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION": "python",
+ "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+ "PYTHONUNBUFFERED": "1",
+ }
+)
+# Safely default to 'tpu' only if JAX_PLATFORMS is not already provided by the Pathways orchestrator
+os.environ.setdefault("JAX_PLATFORMS", "tpu")
+
+
+# --- 2. SETUP PATHS ---
+from maxtext.trainers.post_train.rl.train_rl import rl_train, setup_configs_and_devices
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+login(token=HF_TOKEN)
+
+MODEL_NAME = "llama3.1-8b"
+TOKENIZER_PATH = "meta-llama/Llama-3.1-8B-Instruct"
+
+# Safely grab the native bucket path from Kubernetes, fallback to local if testing
+YOUR_GCS_BUCKET = os.environ.get(
+ "GCS_OUTPUT_PATH", f"{MAXTEXT_PKG_DIR}/fallback_output"
+)
+
+# Pull the base name from K8s, or use a timestamp
+base_name = os.environ.get(
+ "RUN_NAME", datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
+)
+
+# Unconditionally force "v5e-multi" onto the front of it
+RUN_NAME = f"v5e-multi-{base_name}"
+
+# Send the massive converted model and checkpoints directly to the cloud bucket
+MODEL_CHECKPOINT_PATH = f"{YOUR_GCS_BUCKET}/llama_checkpoint"
+
+# MaxText uses `base_output_directory` as the root.
+# It will automatically append `RUN_NAME/checkpoints/` to it.
+OUTPUT_DIRECTORY = YOUR_GCS_BUCKET
+
+CHAT_TEMPLATE_PATH = f"{MAXTEXT_PKG_DIR}/examples/chat_templates/gsm8k_rl.json"
+
+# POINT EXACTLY TO /0/items AS PER THE DEMO NOTEBOOK
+LOAD_PATH = f"{MODEL_CHECKPOINT_PATH}/0/items"
+
+# --- 3. CONVERSION (Runs only if needed) ---
+if not os.path.exists(LOAD_PATH):
+ print("๐ Starting local conversion...")
+
+ # Use subprocess for the conversion
+ conversion_cmd = (
+ f"JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext "
+ f"{MAXTEXT_PKG_DIR}/configs/base.yml "
+ f"model_name={MODEL_NAME} "
+ f"base_output_directory={MODEL_CHECKPOINT_PATH} "
+ f"hf_access_token={HF_TOKEN} "
+ f"use_multimodal=false scan_layers=true skip_jax_distributed_system=True"
+ )
+
+ result = subprocess.run(conversion_cmd, shell=True, executable="/bin/bash")
+ if result.returncode != 0:
+ raise RuntimeError("Conversion failed!")
+else:
+ print(f"โ
Checkpoint already exists at {LOAD_PATH}. Skipping conversion!")
+
+# --- 4. MLFLOW SETUP & LOGGING INTERCEPTOR ---
+# Initialize MLflow strictly on the main thread
+mlflow.set_tracking_uri(
+ os.environ.get("MLFLOW_TRACKING_URI", "http://mlflow-service:5000")
+)
+mlflow.set_experiment("MaxText-RL-GRPO-v5e-multi")
+
+print("๐ Connecting to MLflow database...")
+active_run = mlflow.start_run(run_name=f"Llama3.1-8B-GRPO-{RUN_NAME}")
+MLFLOW_RUN_ID = active_run.info.run_id
+mlflow_client = MlflowClient()
+
+original_write_scalars = clu.metric_writers.MultiWriter.write_scalars
+
+
+def patched_write_scalars(self, step: int, scalars: dict):
+ original_write_scalars(self, step, scalars)
+ mlflow_metrics = {
+ k: float(v)
+ for k, v in scalars.items()
+ if isinstance(v, (jnp.ndarray, float, int))
+ }
+ try:
+ # Pass the entire dictionary at once using the thread-safe client
+ mlflow_client.log_metrics(MLFLOW_RUN_ID, mlflow_metrics, step=int(step))
+ except Exception as e:
+ pass # Silently pass so we don't break the TPU training loop
+
+
+clu.metric_writers.MultiWriter.write_scalars = patched_write_scalars
+
+original_write_texts = clu.metric_writers.MultiWriter.write_texts
+
+
+def patched_write_texts(self, step: int, texts: dict):
+ original_write_texts(self, step, texts)
+ try:
+ # Dynamically find the keys, handling prefixes like "eval/" or "train/"
+ prompt_key = next((k for k in texts.keys() if "prompt" in k.lower()), None)
+ comp_key = next((k for k in texts.keys() if "completion" in k.lower()), None)
+
+ if prompt_key and comp_key:
+ # Tag it visually so you know exactly which phase is printing
+ phase = "๐งช EVALUATION" if "eval" in prompt_key.lower() else "๐ง TRAINING"
+ print(f"\n" + "=" * 20 + f" {phase} STEP {step} SAMPLE " + "=" * 20)
+
+ prompt = texts[prompt_key][0]
+ completion = texts[comp_key][0]
+
+ import numpy as np
+
+ if isinstance(prompt, np.ndarray):
+ prompt = prompt.item() if prompt.size == 1 else str(prompt)
+ if isinstance(completion, np.ndarray):
+ completion = (
+ completion.item() if completion.size == 1 else str(completion)
+ )
+
+ print(f"โ [{prompt_key.upper()}]:\n{prompt}\n")
+ print(f"๐ค [{comp_key.upper()}]:\n{completion}\n")
+ print("=" * 70 + "\n", flush=True)
+ except Exception:
+ pass
+
+
+clu.metric_writers.MultiWriter.write_texts = patched_write_texts
+
+import jax.numpy as jnp
+
+# --- MONKEY PATCHES (For MaxText v0.2.1 / Tunix) ---
+from maxtext.inference.vllm_decode import VllmRollout as MaxText_VllmRollout
+
+try:
+ from tunix.rl.rollout.vllm_rollout import VllmRollout as Tunix_VllmRollout
+except ImportError:
+ Tunix_VllmRollout = None
+
+
+def apply_universal_patches(TargetClass):
+ orig_logps = TargetClass.get_per_token_logps
+
+ def patched_logps(self, *args, **kwargs):
+ mask = kwargs.pop("completion_mask", None)
+ results = orig_logps(self, *args, **kwargs)
+
+ target_len = mask.shape[-1] if mask is not None else 1792
+
+ def pad_sequence(seq):
+ seq_arr = jnp.array(seq)
+ if seq_arr.size == 0:
+ return jnp.zeros(target_len)
+ pad_amount = target_len - seq_arr.shape[0]
+ if pad_amount > 0:
+ return jnp.pad(seq_arr, (0, pad_amount), constant_values=0.0)
+ return seq_arr[:target_len]
+
+ if isinstance(results, list):
+ return jnp.stack([pad_sequence(s) for s in results])
+ elif isinstance(results, dict):
+ return {
+ k: jnp.stack([pad_sequence(s) for s in v]) if isinstance(v, list) else v
+ for k, v in results.items()
+ }
+ return results
+
+ TargetClass.get_per_token_logps = patched_logps
+
+
+apply_universal_patches(MaxText_VllmRollout)
+if Tunix_VllmRollout:
+ apply_universal_patches(Tunix_VllmRollout)
+
+# --- 5. TRAINING CONFIGURATION ---
+config_argv = [
+ "",
+ f"{MAXTEXT_PKG_DIR}/configs/post_train/rl.yml",
+ f"model_name={MODEL_NAME}",
+ f"tokenizer_path={TOKENIZER_PATH}",
+ f"run_name={RUN_NAME}",
+ f"load_parameters_path={LOAD_PATH}",
+ f"base_output_directory={OUTPUT_DIRECTORY}",
+ f"hf_access_token={HF_TOKEN}",
+ f"chat_template_path={CHAT_TEMPLATE_PATH}",
+ f"vllm_hf_config_path={TOKENIZER_PATH}",
+ "rl.loss_algo=grpo",
+ "use_pathways=True",
+ "debug.rl=True",
+ "rl.rollout_engine=vllm",
+ "rollout_tensor_parallelism=8",
+ "rollout_data_parallelism=1",
+ "rl.reasoning_start_token=''",
+ "rl.reasoning_end_token=''",
+ "rl.solution_start_token=''",
+ "rl.solution_end_token=''",
+ # --- BATCHING & MEMORY FIXES ---
+ "batch_size=2", # Down from 4 to save memory
+ "rl.num_generations=8",
+ "max_target_length=1024", # Restored to MaxText's default
+ "hbm_utilization_vllm=0.37", # The v5e "Goldilocks" zone we calculated
+ "num_batches=150", # Quick test run
+ # --- CATASTROPHIC FORGETTING FIXES ---
+ "learning_rate=5e-7", # Much slower than the 3e-6 default
+ "rl.grpo_beta=0.25", # Stronger leash than the 0.08 default
+ "rl.penalty_reward=-0.1", # A gentle nudge instead of a harsh -0.5 punishment
+ # --- FIXED RL PARAMS ---
+ "rl.num_iterations=1",
+ "gradient_clipping_threshold=1.0",
+ "add_eos=True",
+ "log_period=10",
+ "return_log_prob=True",
+ "checkpoint_period=25",
+ "save_checkpoint_on_completion=True",
+ # --- EVALUATION ---
+ "num_test_batches=25",
+ "eval_interval=100",
+ # --- ML DIAGNOSTICS CONFIGURATION ---
+ "managed_mldiagnostics=True", # Enable the managed ML Diagnostics platform
+ "managed_mldiagnostics_run_group=GRPO_RL", # (Optional) Group multiple runs under this category
+ "profiler=xplane", # Enable Google Cloud profiling traces
+ "upload_all_profiler_results=True", # Capture and upload multi-host profiles from all TPU hosts
+]
+
+# --- 6. EXECUTION ---
+trainer_config, sampler_config, trainer_devices, sampler_devices = (
+ setup_configs_and_devices(config_argv)
+)
+
+print(f"๐ฅ Training starting on {len(jax.devices())} TPUs...")
+try:
+ rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices)
+finally:
+ # Ensure the MLflow run is safely closed even if an error occurs
+ mlflow.end_run()
+
+print("๐ Training successfully completed.")
diff --git a/docs/platforms/gke/base/use-cases/reinforcement-learning/multi-host-tpu-grpo/README.md b/docs/platforms/gke/base/use-cases/reinforcement-learning/multi-host-tpu-grpo/README.md
new file mode 100644
index 000000000..e34ef2bd4
--- /dev/null
+++ b/docs/platforms/gke/base/use-cases/reinforcement-learning/multi-host-tpu-grpo/README.md
@@ -0,0 +1,141 @@
+# Multi-host reinforcement learning with TPUs on Google Kubernetes Engine (GKE) using Pathways and JobSet
+
+This example implements distributed multi-host reinforcement learning using
+Group Relative Policy Optimization (GRPO) and MaxText on TPUs on Google
+Kubernetes Engine (GKE).
+
+It integrates **MaxText** (for distributed FSDP model training), **vLLM** (for
+high-throughput rollout generation), and **Tunix** (the RL bridge) on a
+multi-host TPU v5e-16 slice (`v5e-4x4`) orchestrated via **Pathways** and
+**JobSet** to fine-tune Llama-3.1-8B-Instruct.
+
+This example is built on top of the
+[GKE Reinforcement Learning reference architecture](/docs/platforms/gke/base/use-cases/reinforcement-learning/README.md).
+
+## Before you begin
+
+- The
+ [GKE Reinforcement Learning reference implementation](/platforms/gke/base/use-cases/reinforcement-learning/terraform/README.md)
+ is deployed and configured.
+
+- Get access to the model.
+
+ - For Llama-3.1:
+ - Accept the terms of the license on the Hugging Face model page.
+ - [**meta-llama/Llama-3.1-8B-Instruct**](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)
+
+- Ensure your
+ [Hugging Face Hub **Read** access token](/platforms/gke/base/core/huggingface/initialize/README.md)
+ has been added to Secret Manager.
+
+- Hardware & Storage Prerequisites:
+ - **Hardware**: This configuration is tuned for a multi-host **TPU v5e-16**
+ (`v5e-4x4`) slice topology.
+ - **Storage**: GCS bucket (configured via the dataset bucket name) used as a
+ synchronization directory (`pathwaysDir`) for inter-node communication.
+
+## Create and configure the Google Cloud resources
+
+- Deploy the multi-host reinforcement learning on TPU resources.
+
+ ```shell
+ export TF_PLUGIN_CACHE_DIR="${ACP_REPO_DIR}/.terraform.d/plugin-cache"
+ cd ${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host && \
+ rm -rf .terraform/ terraform.tfstate* && \
+ terraform init && \
+ terraform plan -input=false -out=tfplan && \
+ terraform apply -input=false tfplan && \
+ rm tfplan
+ ```
+
+## Build the container images
+
+- Source the environment configuration.
+
+ ```shell
+ source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/scripts/set_environment_variables.sh"
+ ```
+
+- Build the container image for the TPU reinforcement learning trainer.
+
+ ```shell
+ export TF_PLUGIN_CACHE_DIR="${ACP_REPO_DIR}/.terraform.d/plugin-cache"
+ cd ${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host && \
+ rm -rf .terraform/ terraform.tfstate* && \
+ terraform init && \
+ terraform plan -input=false -out=tfplan && \
+ terraform apply -input=false tfplan && \
+ rm tfplan
+ ```
+
+ > The build usually takes 10 to 15 minutes.
+
+## Deploy the reinforcement learning workload
+
+- Source the environment configuration.
+
+ ```shell
+ source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/scripts/set_environment_variables.sh"
+ ```
+
+- Configure the deployment.
+
+ ```shell
+ "${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/configure_job.sh"
+ ```
+
+- Deploy the reinforcement learning workload.
+
+ ```shell
+ kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/v5e-4x4-llama-3-1-8b-instruct"
+ ```
+
+- Watch the reinforcement learning job until it is complete.
+
+ ```shell
+ watch --color --interval 5 --no-title \
+ "kubectl --namespace=${rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name} get pathwaysjob/reinforcement-learning-maxtext-grpo-v5e-4x4-llama-3-1-8b-instruct
+ echo '\nLogs(last 10 lines):'
+ kubectl --namespace=${rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name} logs pathwaysjob/reinforcement-learning-maxtext-grpo-v5e-4x4-llama-3-1-8b-instruct --tail 10"
+ ```
+
+## Viewing Metrics (MLflow & TensorBoard)
+
+MaxText logs step metrics directly to TensorBoard format during execution. The
+`train.py` script automatically packages these logs and attaches them to
+**MLflow** as artifacts upon run completion.
+
+### Accessing the MLflow UI
+
+Because MLflow runs inside the cluster, you can port-forward the service to view
+the dashboard locally:
+
+1. **Port-forward the MLflow Service:**
+
+ ```shell
+ kubectl port-forward --namespace=${rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name} svc/mlflow-service 5000:5000
+ ```
+
+2. **Open your Browser:** Navigate to `http://localhost:5000`
+
+3. **View Experiment Runs:**
+ - Select the `MaxText-RL-GRPO-v5e-multi` experiment.
+ - Click on your active run (e.g., `Llama3.1-8B-GRPO-...`).
+ - Inspect logged metrics (policy loss, reward values, KL divergence) and
+ access attached TensorBoard log archives in the **Artifacts** section.
+
+## Pathways & JobSet Architecture
+
+Because this pipeline spans across multiple hosts, it leverages the
+**PathwaysJob** Custom Resource:
+
+1. **Pathways Orchestration**: The `PathwaysJob` operator creates a highly
+ optimized Pathways cluster consisting of a resource manager (server), proxy
+ server, and worker node pools.
+2. **Underlying JobSet API**: Lifecycle synchronization and reliable process
+ startup across distinct TPU hosts are managed under the hood by GKE's JobSet
+ controller.
+3. **Inter-Host GCS Synced Logging**: Multi-host coordination relies on a shared
+ Google Cloud Storage subdirectory specified in `spec.pathwaysDir`. The
+ workload's service account is granted `roles/storage.objectAdmin` on the
+ dataset bucket to enable transparent file-based handshakes.
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/job.yaml b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/job.yaml
new file mode 100644
index 000000000..f27a5a5e7
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/job.yaml
@@ -0,0 +1,50 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+---
+apiVersion: pathways-job.pathways.domain/v1
+kind: PathwaysJob
+metadata:
+ name: reinforcement-learning-maxtext-grpo
+ namespace: replaced-by-kustomize
+spec:
+ maxRestarts: 0
+ pathwaysDir: replaced-by-kustomize
+ controller:
+ template:
+ spec:
+ restartPolicy: Never
+ containers:
+ - name: grpo-trainer
+ image: replaced-by-kustomize
+ env:
+ - name: JAX_PLATFORMS
+ value: "proxy"
+ - name: JAX_BACKEND_TARGET
+ value: "grpc://127.0.0.1:29000"
+ - name: ENABLE_PATHWAYS_PERSISTENCE
+ value: "1"
+ - name: MLFLOW_TRACKING_URI
+ value: "http://mlflow-service:5000"
+ - name: HF_TOKEN
+ valueFrom:
+ secretKeyRef:
+ name: hf-secret
+ key: token
+ - name: GCS_OUTPUT_PATH
+ value: replaced-by-kustomize
+ serviceAccountName: replaced-by-kustomize
+ workers:
+ - type: replaced-by-kustomize
+ topology: replaced-by-kustomize
+ numSlices: 1
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/kustomization.yaml b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/kustomization.yaml
new file mode 100644
index 000000000..130df03c4
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/kustomization.yaml
@@ -0,0 +1,25 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+---
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+
+configMapGenerator:
+ - envs:
+ - runtime.env
+ name: runtime
+ namespace: replaced-by-kustomize
+
+resources:
+ - job.yaml
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/templates/runtime.tpl.env b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/templates/runtime.tpl.env
new file mode 100644
index 000000000..e4a74afd3
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/base/templates/runtime.tpl.env
@@ -0,0 +1,4 @@
+CONTAINER_IMAGE_URL=${rl_tpu_maxtext_grpo_multi_host_image_url}
+INFERENCE_KUBERNETES_NAMESPACE=${rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name}
+INFERENCE_KUBERNETES_SERVICE_ACCOUNT=${rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name}
+PATHWAYS_DIR=gs://${rl_dataset_bucket_name}/pathways
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/configure_job.sh b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/configure_job.sh
new file mode 100755
index 000000000..684b97f5c
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/configure_job.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+set -o errexit
+set -o nounset
+set -o pipefail
+
+MY_PATH="$(
+ cd "$(dirname "$0")" >/dev/null 2>&1
+ pwd -P
+)"
+
+source "${MY_PATH}/../../terraform/_shared_config/scripts/set_environment_variables.sh"
+
+envsubst < "${MY_PATH}/base/templates/runtime.tpl.env" > "${MY_PATH}/base/runtime.env"
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/v5e-4x4-llama-3-1-8b-instruct/kustomization.yaml b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/v5e-4x4-llama-3-1-8b-instruct/kustomization.yaml
new file mode 100644
index 000000000..2c043f6ef
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/kubernetes-manifests/rl-tpu-maxtext-grpo-multi-host/v5e-4x4-llama-3-1-8b-instruct/kustomization.yaml
@@ -0,0 +1,72 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+---
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+
+nameSuffix: "-v5e-4x4-llama-3-1-8b-instruct"
+
+patches:
+ - patch: |
+ - op: replace
+ path: /spec/workers/0/type
+ value: ct5lp-hightpu-8t
+ - op: replace
+ path: /spec/workers/0/topology
+ value: 4x4
+ target:
+ kind: PathwaysJob
+ name: reinforcement-learning-maxtext-grpo
+
+replacements:
+ - source:
+ fieldPath: data.CONTAINER_IMAGE_URL
+ kind: ConfigMap
+ name: runtime
+ targets:
+ - fieldPaths:
+ - spec.controller.template.spec.containers.[name=grpo-trainer].image
+ select:
+ kind: PathwaysJob
+ - source:
+ fieldPath: data.INFERENCE_KUBERNETES_NAMESPACE
+ kind: ConfigMap
+ name: runtime
+ targets:
+ - fieldPaths:
+ - metadata.namespace
+ select:
+ kind: PathwaysJob
+ - source:
+ fieldPath: data.INFERENCE_KUBERNETES_SERVICE_ACCOUNT
+ kind: ConfigMap
+ name: runtime
+ targets:
+ - fieldPaths:
+ - spec.controller.template.spec.serviceAccountName
+ select:
+ kind: PathwaysJob
+ - source:
+ fieldPath: data.PATHWAYS_DIR
+ kind: ConfigMap
+ name: runtime
+ targets:
+ - fieldPaths:
+ - spec.pathwaysDir
+ - spec.controller.template.spec.containers.[name=grpo-trainer].env.[name=GCS_OUTPUT_PATH].value
+ select:
+ kind: PathwaysJob
+
+resources:
+ - ../base
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/outputs.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/outputs.tf
index 70161c1e6..df8d981e0 100644
--- a/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/outputs.tf
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/outputs.tf
@@ -43,3 +43,7 @@ output "rl_tpu_maxtext_grpo_single_host_kubernetes_namespace_name" {
output "rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name" {
value = local.rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name
}
+
+output "rl_tpu_maxtext_grpo_multi_host_image_url" {
+ value = local.rl_tpu_maxtext_grpo_multi_host_image_url
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/reinforcement_learning_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/reinforcement_learning_variables.tf
index c48442241..ff7d74f36 100644
--- a/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/reinforcement_learning_variables.tf
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/_shared_config/reinforcement_learning_variables.tf
@@ -16,6 +16,7 @@ locals {
rl_project_id = var.rl_project_id != null ? var.rl_project_id : var.platform_default_project_id
rl_dataset_bucket_name = var.rl_dataset_bucket_name != null ? var.rl_dataset_bucket_name : "${local.rl_project_id}-${local.unique_identifier_prefix}-dataset"
rl_mlflow_data_bucket_name = var.rl_mlflow_data_bucket_name != null ? var.rl_mlflow_data_bucket_name : "${local.rl_project_id}-${local.unique_identifier_prefix}-mlflow-data"
+ rl_project_id = var.rl_project_id != null ? var.rl_project_id : var.platform_default_project_id
rl_cpu_maxtext_checkpoint_converter_kubernetes_namespace_name = var.rl_cpu_maxtext_checkpoint_converter_kubernetes_namespace_name != null ? var.rl_cpu_maxtext_checkpoint_converter_kubernetes_namespace_name : "${local.unique_identifier_prefix}-checkpoint-converter"
rl_cpu_maxtext_checkpoint_converter_kubernetes_service_account_name = var.rl_cpu_maxtext_checkpoint_converter_kubernetes_service_account_name != null ? var.rl_cpu_maxtext_checkpoint_converter_kubernetes_service_account_name : "${local.unique_identifier_prefix}-checkpoint-converter-sa"
@@ -26,6 +27,10 @@ locals {
rl_tpu_maxtext_grpo_single_host_image_url = var.rl_tpu_maxtext_grpo_single_host_image_url != null ? var.rl_tpu_maxtext_grpo_single_host_image_url : "${local.cloudbuild_ar_image_repository_url}/reinforcement-learning/grpo-single-host:latest"
rl_tpu_maxtext_grpo_single_host_kubernetes_namespace_name = var.rl_tpu_maxtext_grpo_single_host_kubernetes_namespace_name != null ? var.rl_tpu_maxtext_grpo_single_host_kubernetes_namespace_name : "${local.unique_identifier_prefix}-grpo-single-host"
rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name = var.rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name != null ? var.rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name : "${local.unique_identifier_prefix}-grpo-single-host-sa"
+
+ rl_tpu_maxtext_grpo_multi_host_image_url = var.rl_tpu_maxtext_grpo_multi_host_image_url != null ? var.rl_tpu_maxtext_grpo_multi_host_image_url : "${local.cloudbuild_ar_image_repository_url}/reinforcement-learning/grpo-multi-host:latest"
+ rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name = var.rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name != null ? var.rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name : "${local.unique_identifier_prefix}-grpo-multi-host"
+ rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name = var.rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name != null ? var.rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name : "${local.unique_identifier_prefix}-grpo-multi-host-sa"
}
variable "rl_cpu_maxtext_checkpoint_converter_kubernetes_namespace_name" {
@@ -87,3 +92,21 @@ variable "rl_tpu_maxtext_grpo_single_host_kubernetes_service_account_name" {
description = "The Kubernetes service account name for the RL on TPU deployment."
type = string
}
+
+variable "rl_tpu_maxtext_grpo_multi_host_image_url" {
+ default = null
+ description = "The URL for the RL on TPU multi-host container image."
+ type = string
+}
+
+variable "rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name" {
+ default = null
+ description = "The Kubernetes namespace name for the RL on TPU multi-host deployment."
+ type = string
+}
+
+variable "rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name" {
+ default = null
+ description = "The Kubernetes service account name for the RL on TPU multi-host deployment."
+ type = string
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/deploy-standard.sh b/platforms/gke/base/use-cases/reinforcement-learning/terraform/deploy-standard.sh
index 67a7a94cc..cc6395ab1 100755
--- a/platforms/gke/base/use-cases/reinforcement-learning/terraform/deploy-standard.sh
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/deploy-standard.sh
@@ -66,6 +66,7 @@ source "${ACP_PLATFORM_USE_CASE_DIR}/terraform/_shared_config/scripts/set_enviro
declare -a use_case_terraservices=(
"initialize"
"rl_tpu_maxtext_grpo_single_host"
+ "rl_tpu_maxtext_grpo_multi_host"
)
for terraservice in "${use_case_terraservices[@]}"; do
cd "${ACP_PLATFORM_USE_CASE_DIR}/terraform/${terraservice}" &&
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars
new file mode 120000
index 000000000..238bf8e95
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars
@@ -0,0 +1 @@
+../../../_shared_config/_cloudbuild.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf
new file mode 120000
index 000000000..8fade6147
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf
@@ -0,0 +1 @@
+../../../_shared_config/_cloudbuild_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars
new file mode 120000
index 000000000..c9c406bba
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars
@@ -0,0 +1 @@
+../../../_shared_config/_platform.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf
new file mode 120000
index 000000000..7ec64070d
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf
@@ -0,0 +1 @@
+../../../_shared_config/_platform_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars
new file mode 120000
index 000000000..171a27a35
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars
@@ -0,0 +1 @@
+../../../_shared_config/reinforcement_learning.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf
new file mode 120000
index 000000000..79960dd37
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf
@@ -0,0 +1 @@
+../../../_shared_config/reinforcement_learning_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/cloudbuild.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/cloudbuild.tf
new file mode 100644
index 000000000..34d409dbd
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/cloudbuild.tf
@@ -0,0 +1,47 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+locals {
+ image_destination = local.rl_tpu_maxtext_grpo_multi_host_image_url
+}
+
+resource "terraform_data" "submit_docker_build" {
+ input = {
+ acp_root = local.acp_root
+ cloudbuild_project_id = local.cloudbuild_project_id
+ cloudbuild_service_account_id = local.cloudbuild_service_account_id
+ cloudbuild_source_bucket_name = local.cloudbuild_source_bucket_name
+ image_destination = local.image_destination
+ }
+
+ provisioner "local-exec" {
+ command = <<-EOT
+gcloud builds submit \
+--config="container-images/tpu/rl-tpu-maxtext-grpo-multi-host/cloudbuild.yaml" \
+--gcs-source-staging-dir="gs://${self.input.cloudbuild_source_bucket_name}/source" \
+--project="${self.input.cloudbuild_project_id}" \
+--quiet \
+--service-account="${self.input.cloudbuild_service_account_id}" \
+--substitutions=_DESTINATION="${self.input.image_destination}"
+EOT
+ interpreter = ["bash", "-c"]
+ working_dir = self.input.acp_root
+ }
+
+ triggers_replace = {
+ cloudbuild_yaml_hash = filebase64sha256("${local.acp_root}/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/cloudbuild.yaml")
+ dockerfile_hash = filebase64sha256("${local.acp_root}/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/Dockerfile")
+ source_hash = sha256(join("", [for file in fileset("${local.acp_root}/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src", "**") : filesha256("${local.acp_root}/container-images/tpu/rl-tpu-maxtext-grpo-multi-host/src/${file}")]))
+ }
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/local_file.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/local_file.tf
new file mode 100644
index 000000000..ef13fc1bb
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/local_file.tf
@@ -0,0 +1,17 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+locals {
+ acp_root = "${path.module}/../../../../../../../../.."
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/versions.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/versions.tf
new file mode 100644
index 000000000..3beb7daeb
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/images/tpu/rl_tpu_maxtext_grpo_multi_host/versions.tf
@@ -0,0 +1,32 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+terraform {
+ required_version = ">= 1.5.7"
+
+ required_providers {
+ google = {
+ source = "hashicorp/google"
+ version = "6.49.2"
+ }
+ local = {
+ source = "hashicorp/local"
+ version = "2.5.3"
+ }
+ }
+
+ provider_meta "google" {
+ module_name = "cloud-solutions/acp_rl_images_tpu_rl_tpu_maxtext_grpo_multi_host_deploy-v1"
+ }
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars
new file mode 120000
index 000000000..2af7bbaaa
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild.auto.tfvars
@@ -0,0 +1 @@
+../_shared_config/_cloudbuild.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf
new file mode 120000
index 000000000..dd199215c
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cloudbuild_variables.tf
@@ -0,0 +1 @@
+../_shared_config/_cloudbuild_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster.auto.tfvars
new file mode 120000
index 000000000..04c4ae417
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster.auto.tfvars
@@ -0,0 +1 @@
+../_shared_config/_cluster.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster_variables.tf
new file mode 120000
index 000000000..6713167a1
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_cluster_variables.tf
@@ -0,0 +1 @@
+../_shared_config/_cluster_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars
new file mode 120000
index 000000000..f898b3b5a
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform.auto.tfvars
@@ -0,0 +1 @@
+../_shared_config/_platform.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf
new file mode 120000
index 000000000..f928d86dd
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_platform_variables.tf
@@ -0,0 +1 @@
+../_shared_config/_platform_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars
new file mode 120000
index 000000000..f56697856
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning.auto.tfvars
@@ -0,0 +1 @@
+../_shared_config/reinforcement_learning.auto.tfvars
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf
new file mode 120000
index 000000000..f7d4bb73a
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/_reinforcement_learning_variables.tf
@@ -0,0 +1 @@
+../_shared_config/reinforcement_learning_variables.tf
\ No newline at end of file
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/iam.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/iam.tf
new file mode 100644
index 000000000..c1b8aa347
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/iam.tf
@@ -0,0 +1,26 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the_unique_identifier_prefix.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+locals {
+ cluster_wi_principal_prefix = "principal://iam.googleapis.com/projects/${data.google_project.cluster.number}/locations/global/workloadIdentityPools/${data.google_project.cluster.project_id}.svc.id.goog/subject"
+ rl_cpu_mlflow_ksa_member = "${local.cluster_wi_principal_prefix}/ns/${local.rl_cpu_mlflow_kubernetes_namespace_name}/sa/${local.rl_cpu_mlflow_kubernetes_service_account_name}"
+ rl_tpu_maxtext_grpo_multi_host_ksa_member = "${local.cluster_wi_principal_prefix}/ns/${local.rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name}/sa/${local.rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name}"
+}
+
+resource "google_storage_bucket_iam_member" "dataset_bucket_multi_host_storage_object_admin" {
+ bucket = google_storage_bucket.dataset.name
+ member = local.rl_tpu_maxtext_grpo_multi_host_ksa_member
+ role = "roles/storage.objectAdmin"
+}
diff --git a/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/kubernetes.tf b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/kubernetes.tf
new file mode 100644
index 000000000..153b5f46b
--- /dev/null
+++ b/platforms/gke/base/use-cases/reinforcement-learning/terraform/rl_tpu_maxtext_grpo_multi_host/kubernetes.tf
@@ -0,0 +1,117 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+locals {
+ kubeconfig_directory = "${path.module}/../../../../kubernetes/kubeconfig"
+ kubeconfig_file = "${local.kubeconfig_directory}/${local.kubeconfig_file_name}"
+
+ namespaces_directory = "${local.manifests_directory_root}/namespaces"
+
+ workloads = {
+ rl_cpu_mlflow = {
+ directory = "${local.namespaces_directory}/${local.rl_cpu_mlflow_kubernetes_namespace_name}"
+ namespace = local.rl_cpu_mlflow_kubernetes_namespace_name
+ service_account = local.rl_cpu_mlflow_kubernetes_service_account_name
+ }
+ rl_tpu_maxtext_grpo_multi_host = {
+ directory = "${local.namespaces_directory}/${local.rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name}"
+ namespace = local.rl_tpu_maxtext_grpo_multi_host_kubernetes_namespace_name
+ service_account = local.rl_tpu_maxtext_grpo_multi_host_kubernetes_service_account_name
+ }
+ }
+}
+
+data "local_file" "kubeconfig" {
+ filename = local.kubeconfig_file
+}
+
+resource "terraform_data" "namespaces" {
+ for_each = local.workloads
+
+ input = {
+ directory = each.value.directory
+ namespace = each.value.namespace
+ service_account = each.value.service_account
+ }
+
+ provisioner "local-exec" {
+ command = <