diff --git a/src/maxtext/training_engine/abstract_engine.py b/src/maxtext/training_engine/abstract_engine.py
new file mode 100644
index 0000000000..8207e326a9
--- /dev/null
+++ b/src/maxtext/training_engine/abstract_engine.py
@@ -0,0 +1,248 @@
+# 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.
+
+"""Trainer abstractions.
+
+Defines the core Trainer interface, data payload interfaces, and on-device
+metrics structures (WeightedMetric, MetricsBuffer) used by the training loop.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Callable
+import dataclasses
+from typing import Any
+
+import flax.struct
+import jax
+from jax.typing import ArrayLike  # pylint: disable=g-importing-member
+
+
+@flax.struct.dataclass
+class WeightedMetric:
+  """A metric that requires weighted reduction.
+
+  Attributes:
+    unreduced_sum: Sum of the metric values across tokens/examples.
+    denominator: Weight or count of valid tokens/examples.
+    eps: Optional epsilon added to denominator for numerical stability.
+    min_denom: Optional minimum bound for the denominator.
+  """
+
+  unreduced_sum: jax.Array
+  denominator: jax.Array
+  eps: float | None = flax.struct.field(default=None, pytree_node=False)
+  min_denom: float | None = flax.struct.field(default=None, pytree_node=False)
+
+  def compute_scale(self) -> jax.Array:
+    """Safely computes the scale factor (1 / denominator) with bounds.
+
+    Returns:
+      Safe scaling factor array preventing division-by-zero NaNs.
+    """
+    denom = self.denominator
+    if self.min_denom is not None:
+      denom = jax.numpy.maximum(denom, self.min_denom)
+    if self.eps is not None:
+      denom = denom + self.eps
+    safe_denom = jax.numpy.where(denom == 0, 1.0, denom)
+    scale = 1.0 / safe_denom
+    return jax.numpy.where(denom == 0, 0.0, scale)
+
+  def compute(self) -> jax.Array:
+    """Safely computes total / count with numerical stability bounds.
+
+    Returns:
+      Reduced metric array equal to unreduced_sum * compute_scale().
+    """
+    return self.unreduced_sum * self.compute_scale()
+
+
+@flax.struct.dataclass
+class MetricsBuffer:
+  """A buffer for storing and aggregating unreduced metrics on-device.
+
+  Attributes:
+    id: Identifier for the buffer (e.g., training iteration or step index).
+    weighted_metrics: Dictionary of WeightedMetric objects on accelerator HBM.
+    scalar_metrics: Dictionary of scalar JAX arrays on accelerator HBM.
+    aggregation_fns: Host-side reduction/aggregation callbacks (untraced).
+    mode: Execution mode string ("train" or "eval").
+  """
+
+  id: Any
+  weighted_metrics: dict[str, WeightedMetric] = flax.struct.field(
+      default_factory=dict
+  )
+  scalar_metrics: dict[str, jax.Array] = flax.struct.field(default_factory=dict)
+  aggregation_fns: dict[str, Callable[[jax.Array], Any]] = flax.struct.field(
+      default_factory=dict, pytree_node=False
+  )
+  mode: str = flax.struct.field(default="train", pytree_node=False)
+
+
+@dataclasses.dataclass(kw_only=True)
+class TrainerPayload(abc.ABC):
+  """Base class for packed micro-batches ready for gradient descent.
+
+  The base carries only what generic machinery must read to stay
+  algorithm-agnostic. Algorithm-specific tensors live on subclasses and are
+  reached by the trainer's gen_model_input_fn, not by the generic loop. Users
+  subclass this to carry their own fields.
+
+  Attributes:
+    token_ids: [B, T] token IDs. By default, structured as left-padded prompt
+      tokens concatenated with right-padded completion tokens.
+    token_mask: [B, T] token mask to differentiate padding tokens from valid
+      tokens.
+    segment_ids: Optional [B, T] packing segment ids.
+  """
+
+  token_ids: ArrayLike
+  token_mask: ArrayLike
+  segment_ids: ArrayLike | None = None
+
+
+@dataclasses.dataclass
+class TrainingConfig:
+  """Configuration for the abstract trainer.
+
+  Defines standard hyperparameters and operational settings for the ML training
+  loop.
+  """
+
+  eval_every_n_steps: int = 0
+  max_steps: int | None = None
+  gradient_accumulation_steps: int | None = None
+  checkpoint_root_directory: str | None = None
+  metrics_prefix: str = ""
+  max_inflight_computations: int = 2
+
+
+class AbstractTrainingEngine(abc.ABC):
+  """Core trainer interface executing model updates and Multi-Tier Checkpointing.
+
+  The Trainer owns the model weights in accelerator HBM and executes forward/
+  backward passes, weight updates, evaluation steps, and checkpoint saving/
+  restoring.
+  """
+
+  @abc.abstractmethod
+  def __init__(self, training_config: TrainingConfig) -> None:
+    """Initializes the Trainer based on the training configuration.
+
+    Args:
+      training_config: Training hyperparameters and runtime configuration.
+    """
+
+  @abc.abstractmethod
+  def with_loss_fn(self, customized_fn: Callable[..., Any]) -> None:
+    """Updates the trainer's loss function.
+
+    Args:
+      customized_fn: Custom loss function callable.
+    """
+
+  @abc.abstractmethod
+  def with_gen_model_input_fn(
+      self, gen_model_input_fn: Callable[[Any], dict[str, Any]]
+  ) -> "AbstractTrainingEngine":
+    """Sets the last-mile adapter mapping a payload to the loss fn's kwargs.
+
+    This adapter enables the trainer to accept arbitrary payloads (SFT, RL,
+    etc.) by transforming them into kwargs for the loss function via
+    `gen_model_input_fn(payload)`.
+    Args:
+      gen_model_input_fn: Maps a payload to a dict of loss-fn keyword arguments.
+
+    Returns:
+      self, for chaining.
+    """
+
+  @abc.abstractmethod
+  def compile(self, dummy_data: TrainerPayload) -> None:
+    """Triggers JAX compilation. `with_loss_fn` must be called first.
+
+    Args:
+      dummy_data: Payload with representative shapes used for JAX tracing.
+    """
+
+  @abc.abstractmethod
+  def fwd_bwd(self, payload: TrainerPayload) -> None:
+    """Executes forward and backward passes.
+
+    Metrics are cached to overlap train steps.
+
+    Args:
+      payload: Packed micro-batch payload for training.
+    """
+
+  @abc.abstractmethod
+  def update(self) -> None:
+    """Executes a model weight update step using accumulated gradients."""
+
+  @abc.abstractmethod
+  def eval_step(self, payload: TrainerPayload, **kwargs: Any) -> None:
+    """Executes one evaluation step on the given payload.
+
+    Args:
+      payload: Packed micro-batch payload for evaluation.
+      **kwargs: Additional evaluation keyword arguments.
+    """
+
+  @abc.abstractmethod
+  def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None:
+    """Forces the trainer to serialize its state (model + optimizer).
+
+    Args:
+      metadata: Checkpoint identifier or UUID metadata pytree.
+      **kwargs: Additional checkpointing keyword arguments.
+    """
+
+  @abc.abstractmethod
+  def restore_checkpoint(self, **kwargs: Any) -> Any:
+    """Restores state from latest checkpoint and returns the metadata pytree.
+
+    The returned metadata (e.g., global_step) matches what was stored in
+    save_checkpoint.
+
+    Args:
+      **kwargs: Additional restoration keyword arguments.
+
+    Returns:
+      The metadata PyTree stored with the checkpoint.
+    """
+
+  @abc.abstractmethod
+  def get_metrics(self, clear_cache: bool = True) -> MetricsBuffer:
+    """Returns cached metrics and optionally clears the metrics cache.
+
+    Args:
+      clear_cache: Whether to reset cached metrics after retrieval.
+
+    Returns:
+      The accumulated on-device MetricsBuffer.
+    """
+
+  @abc.abstractmethod
+  def prepare_weight_sync(self, **kwargs: Any) -> Any:
+    """Stages weights for transfer and returns metadata/coordinates.
+
+    Args:
+      **kwargs: Weight staging configuration parameters.
+
+    Returns:
+      Synchronization endpoints or file coordinates for weight transfer.
+    """
diff --git a/src/maxtext/training_engine/abstract_engine_test.py b/src/maxtext/training_engine/abstract_engine_test.py
new file mode 100644
index 0000000000..5984d8ad78
--- /dev/null
+++ b/src/maxtext/training_engine/abstract_engine_test.py
@@ -0,0 +1,113 @@
+# 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.
+
+"""Tests for trainer abstractions."""
+
+import dataclasses
+from typing import Any
+from absl.testing import absltest
+import jax.numpy as jnp
+from maxtext.training_engine import abstract_engine
+
+
+@dataclasses.dataclass(kw_only=True)
+class DummyPayload(abstract_engine.TrainerPayload):
+  """Dummy payload for testing."""
+
+  data: str = "batch_1"
+
+
+class DummyTrainingEngine(abstract_engine.AbstractTrainingEngine):
+  """Minimal concrete implementation of AbstractTrainingEngine for unit testing."""
+
+  def __init__(self, training_config: abstract_engine.TrainingConfig) -> None:
+    self.config = training_config
+    self.loss_fn = None
+    self.compiled = False
+    self.fwd_bwd_called = 0
+    self.update_called = 0
+    self.checkpoint_saved_with_metadata: Any = None
+    self.restored_metadata: Any = {"global_step": 42}
+
+  def with_loss_fn(self, customized_fn: Any) -> None:
+    self.loss_fn = customized_fn
+
+  def with_gen_model_input_fn(
+      self, gen_model_input_fn: Any
+  ) -> "DummyTrainingEngine":
+    self._gen_model_input_fn = gen_model_input_fn
+    return self
+
+  def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None:
+    self.compiled = True
+
+  def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None:
+    self.fwd_bwd_called += 1
+
+  def update(self) -> None:
+    self.update_called += 1
+
+  def eval_step(
+      self, payload: abstract_engine.TrainerPayload, **kwargs: Any
+  ) -> None:
+    pass
+
+  def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None:
+    self.checkpoint_saved_with_metadata = metadata
+
+  def restore_checkpoint(self, **kwargs: Any) -> Any:
+    return self.restored_metadata
+
+  def get_metrics(
+      self, clear_cache: bool = True
+  ) -> abstract_engine.MetricsBuffer:
+    return abstract_engine.MetricsBuffer(id=1)
+
+  def prepare_weight_sync(self, **kwargs: Any) -> Any:
+    return {"endpoint": "grpc://dummy-trainer:55555"}
+
+
+class AbstractTrainingEngineTest(absltest.TestCase):
+
+  def test_abstract_training_engine_cannot_be_instantiated_directly(self):
+    config = abstract_engine.TrainingConfig()
+    with self.assertRaises(TypeError):
+      abstract_engine.AbstractTrainingEngine(config)  # pytype: disable=not-instantiable
+
+  def test_dummy_training_engine_implements_abstract_interface(self):
+    config = abstract_engine.TrainingConfig(max_steps=100)
+    t = DummyTrainingEngine(config)
+    self.assertEqual(t.config.max_steps, 100)
+    payload = DummyPayload(
+        data="dummy",
+        token_ids=jnp.ones((2, 2)),
+        token_mask=jnp.ones((2, 2)),
+    )
+    t.compile(payload)
+    self.assertTrue(t.compiled)
+    t.fwd_bwd(payload)
+    self.assertEqual(t.fwd_bwd_called, 1)
+    t.update()
+    self.assertEqual(t.update_called, 1)
+
+  def test_weighted_metric_compute(self):
+    m = abstract_engine.WeightedMetric(
+        unreduced_sum=jnp.array(10.0),
+        denominator=jnp.array(2.0),
+    )
+    self.assertAlmostEqual(float(m.compute()), 5.0)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/src/maxtext/training_engine/checkpointing.py b/src/maxtext/training_engine/checkpointing.py
new file mode 100644
index 0000000000..eca025299c
--- /dev/null
+++ b/src/maxtext/training_engine/checkpointing.py
@@ -0,0 +1,175 @@
+# 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.
+
+"""Checkpointing utilities for MaxText training engine."""
+
+from typing import Any
+from absl import logging
+from flax import nnx
+import jax
+from maxtext.configs import pyconfig
+import orbax.checkpoint as ocp
+
+
+class CheckpointManager:
+  """CheckpointManager wrapper for MaxText training engine."""
+
+  def __init__(
+      self,
+      checkpoint_dir: str,
+      config: pyconfig.HyperParameters,
+  ) -> None:
+    """Initializes the CheckpointManager.
+
+    Args:
+      checkpoint_dir: The root directory for saving checkpoints.
+      config: The training configuration.
+    """
+    self._checkpoint_manager: ocp.CheckpointManager | None = None
+    if checkpoint_dir:
+      self._checkpoint_manager = ocp.CheckpointManager(
+          directory=checkpoint_dir,
+          options=ocp.CheckpointManagerOptions(
+              save_interval_steps=getattr(config, "checkpoint_period", 1),
+              max_to_keep=getattr(config, "max_num_checkpoints_to_keep", None),
+              enable_async_checkpointing=getattr(
+                  config, "async_checkpointing", True
+              ),
+          ),
+      )
+
+  def get_latest_step(self) -> int | None:
+    """Returns the latest checkpoint step."""
+    if self._checkpoint_manager:
+      return self._checkpoint_manager.latest_step()
+    return None
+
+  def save_checkpoint(
+      self,
+      step: int,
+      model: nnx.Module,
+      optimizer: nnx.optimizer.Optimizer | None,
+      custom_metadata: Any,
+  ) -> bool:
+    """Saves the params for the given step.
+
+    Args:
+      step: The step to save the params for.
+      model: The model to save.
+      optimizer: The optimizer to save.
+      custom_metadata: Custom metadata to save with the checkpoint.
+
+    Returns:
+      Whether the checkpoint was saved.
+    """
+    if self._checkpoint_manager is None:
+      logging.info("Checkpointing is disabled, skipping save.")
+      return False
+
+    # Check if the checkpoint already exists at the current step.
+    if self.get_latest_step() == step:
+      logging.info(
+          "Checkpoint already saved at step %d, skipping save.",
+          step,
+      )
+      return False
+
+    params = nnx.state(model)
+    jax.block_until_ready(params)
+    model_cp_args = ocp.args.PyTreeSave(
+        item=params,
+        save_args=jax.tree.map(lambda _: ocp.SaveArgs(), params),
+    )
+    save_args = {"model_params": model_cp_args}
+    if optimizer is not None:
+      optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
+      jax.block_until_ready(optimizer_state)
+      optimizer_cp_args = ocp.args.PyTreeSave(
+          item=optimizer_state,
+          save_args=jax.tree.map(lambda _: ocp.SaveArgs(), optimizer_state),
+      )
+      save_args["optimizer_state"] = optimizer_cp_args
+
+    return self._checkpoint_manager.save(
+        step=step,
+        args=ocp.args.Composite(**save_args),
+        custom_metadata=custom_metadata,
+    )
+
+  def restore_checkpoint(
+      self,
+      model: nnx.Module,
+      optimizer: nnx.optimizer.Optimizer | None,
+      step: int | None = None,
+  ) -> Any:
+    """Restores items from the checkpoint at the given step.
+
+    Args:
+      model: The model to restore the params to.
+      optimizer: The optimizer to restore the state to.
+      step: Optional step index to restore from.
+
+    Returns:
+      The metadata of the restored checkpoint, or None if no checkpoint was
+      restored.
+    """
+    if self._checkpoint_manager is None:
+      logging.info("Checkpointing is disabled, skipping restore.")
+      return None, {}
+
+    if step is None:
+      step = self.get_latest_step()
+      if step is None:
+        logging.info("No checkpoint found, skipping restore.")
+        return None, {}
+
+    metadata = self._checkpoint_manager.metadata(step)
+
+    abstract_params = nnx.state(model)
+    model_args = ocp.args.PyTreeRestore(
+        item=abstract_params,
+        restore_args=ocp.checkpoint_utils.construct_restore_args(
+            target=abstract_params
+        ),
+    )
+    if optimizer is not None and "optimizer_state" in metadata.item_metadata:
+      optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
+      optimizer_args = ocp.args.PyTreeRestore(
+          item=optimizer_state,
+          restore_args=ocp.checkpoint_utils.construct_restore_args(
+              target=optimizer_state
+          ),
+      )
+      restore_args = {
+          "model_params": model_args,
+          "optimizer_state": optimizer_args,
+      }
+    else:
+      restore_args = {"model_params": model_args}
+
+    restored_items = self._checkpoint_manager.restore(
+        step=step,
+        args=ocp.args.Composite(**restore_args),
+    )
+    if "model_params" in restored_items:
+      nnx.update(model, restored_items["model_params"])
+    if optimizer is not None and "optimizer_state" in restored_items:
+      nnx.update(optimizer, restored_items["optimizer_state"])
+
+    return step, metadata
+
+  def close(self) -> None:
+    """Closes the checkpoint manager."""
+    if self._checkpoint_manager:
+      self._checkpoint_manager.close()
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
new file mode 100644
index 0000000000..a44625025b
--- /dev/null
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -0,0 +1,342 @@
+# 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.
+
+"""MaxText concrete implementation of AbstractTrainer for RL post-training.
+
+Adapts MaxText's single-step compilation and execution primitives to implement
+the MaxRL AbstractTrainer interface without running an outer loop.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from absl import logging
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from maxtext.common import common_types
+from maxtext.configs import pyconfig
+from maxtext.trainers.pre_train import train as maxtext_train
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import checkpointing
+from maxtext.training_engine import metrics
+from maxtext.utils import gradient_accumulation
+from maxtext.utils import max_utils
+from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
+from maxtext.utils import train_utils
+
+
+class MaxTextTrainingEngine(abstract_engine.AbstractTrainingEngine):
+  """Concrete trainer wrapping MaxText single-step SPMD execution for NNX models."""
+
+  def __init__(
+      self,
+      training_config: pyconfig.HyperParameters,
+      mesh: jax.sharding.Mesh | None = None,
+  ) -> None:
+    """Initializes the MaxText trainer state and sharded model.
+
+    Args:
+      training_config: MaxText HyperParameters configuration instance.
+      mesh: Optional SPMD device mesh.
+
+    Raises:
+      TypeError: If training_config is not a pyconfig.HyperParameters instance.
+      ValueError: If training_config.model_name is not specified or empty.
+    """
+    if not isinstance(training_config, pyconfig.HyperParameters):
+      raise TypeError(
+          "MaxTextTrainingEngine requires a pyconfig.HyperParameters instance,"
+          f" got {type(training_config).__name__}"
+      )
+    self._config = training_config
+    self._mesh = mesh
+    self._init_rng = jax.random.PRNGKey(
+        getattr(training_config, "init_weights_seed", 0)
+    )
+    self._loss_fn: Callable[..., Any] | None = None
+    self._gen_model_input_fn: Callable[[Any], dict[str, Any]] | None = None
+    self._compiled = False
+    if not getattr(training_config, "model_name", None):
+      raise ValueError("training_config.model_name must be specified")
+    self._model = model_creation_utils.from_pretrained(
+        config=self._config,
+        mesh=self._mesh,
+        model_mode=common_types.MODEL_MODE_TRAIN,
+        rng_key=self._init_rng,
+    )
+    self._state: Any = None
+    self._accumulated_grads: Any = None
+    self._micro_step_count = 0
+    self._cached_losses: list[jax.Array] = []
+    self._learning_rate_schedule, self._optimizer = (
+        train_utils.create_training_optimizer(self._config, self._model)
+    )
+    self._train_step: int = 0
+
+    self._checkpoint_manager = checkpointing.CheckpointManager(
+        checkpoint_dir=getattr(self._config, "checkpoint_directory", ""),
+        config=self._config,
+    )
+    self._metrics_logger = metrics.MetricsLogger(self._config)
+
+  @property
+  def model(self) -> Any:
+    """Returns the NNX model instance."""
+    return self._model
+
+  @model.setter
+  def model(self, new_model: Any) -> None:
+    """Sets the NNX model instance."""
+    self._model = new_model
+
+  @property
+  def optimizer(self) -> Any:
+    """Returns the NNX optimizer instance."""
+    return self._optimizer
+
+  @optimizer.setter
+  def optimizer(self, new_optimizer: Any) -> None:
+    """Sets the NNX optimizer instance."""
+    self._optimizer = new_optimizer
+
+  @property
+  def train_step(self) -> int:
+    """Returns the current step integer."""
+    return self._train_step
+
+  @train_step.setter
+  def train_step(self, step: int) -> None:
+    """Sets the current step integer."""
+    self._train_step = step
+
+  def with_loss_fn(self, customized_fn: Callable[..., Any]) -> None:
+    """Overrides the default autoregressive loss function with a custom RL loss.
+
+    Args:
+      customized_fn: Custom loss callable matching the MaxText loss signature.
+    """
+    self._loss_fn = customized_fn
+    self._compiled = False
+
+  def with_gen_model_input_fn(
+      self, gen_model_input_fn: Callable[[Any], dict[str, Any]]
+  ) -> "MaxTextTrainingEngine":
+    """Sets the last-mile adapter mapping a payload to the loss fn's kwargs."""
+    self._gen_model_input_fn = gen_model_input_fn
+    return self
+
+  def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None:
+    """Triggers SPMD JIT compilation of fwd_bwd, update, and eval steps.
+
+    Args:
+      dummy_data: Sample TrainerPayload providing representative tensor shapes.
+    """
+    self._compiled = True
+
+  def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None:
+    """Executes a micro-batch forward-backward pass and accumulates gradients.
+
+    Args:
+      payload: Packed micro-batch training input.
+    """
+    if self._gen_model_input_fn is not None:
+      batch = self._gen_model_input_fn(payload)
+    else:
+      batch = payload
+    loss_callable = (
+        self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn
+    )
+
+    model = (
+        getattr(self._state, "model", None)
+        if self._state is not None
+        else self._model
+    )
+    if not isinstance(model, nnx.Module):
+      raise TypeError(
+          "MaxRL requires an NNX model (flax.nnx.Module), got"
+          f" {type(model).__name__}"
+      )
+    # TODO(mazumdera): This function call should be pre-compiled.
+    loss, _, micro_grads = (
+        gradient_accumulation.gradient_accumulation_loss_and_grad(
+            loss_callable,
+            self._config,
+            model,
+            None,
+            None,
+            batch,
+            None,
+        )
+    )
+
+    if isinstance(loss, abstract_engine.WeightedMetric):
+      self.record_metrics("loss", loss)
+
+    self._cached_losses.append(loss)
+    if self._accumulated_grads is None:
+      self._accumulated_grads = micro_grads
+    else:
+      self._accumulated_grads = jax.tree.map(
+          jnp.add, self._accumulated_grads, micro_grads
+      )
+    self._micro_step_count += 1
+
+  def update(self) -> None:
+    """Applies accumulated gradients to update NNX model weights in HBM.
+
+    Reuses NNX optimizer step from train.py (lines 511-535).
+    """
+    if self._accumulated_grads is None:
+      return
+    # TODO(mazumdera): The logic below should be pre-compiled.
+    if self._state is not None:
+      # TODO(mazumdera): Figure out how exactly we should normalize the losses
+      # (if at all). Given that inputs are varying in size, it not correct to
+      # simply divide by the number of micro-steps.
+      grads = jax.tree.map(
+          lambda g: g / max(self._micro_step_count, 1), self._accumulated_grads
+      )
+      if getattr(self._config, "gradient_clipping_threshold", 0.0) > 0:
+        grads = maxtext_utils.apply_gradient_clipping(
+            grads, None, self._config.gradient_clipping_threshold
+        )
+      if hasattr(self._state, "apply_gradients"):
+        if getattr(self._config, "skip_step_on_spikes", False):
+          grad_norm = max_utils.l2norm_pytree(grads)
+          mean_loss = (
+              jnp.mean(jnp.array(self._cached_losses))
+              if self._cached_losses
+              else jnp.array(0.0)
+          )
+          self._state.apply_gradients(
+              grads, loss=mean_loss, grad_norm=grad_norm
+          )
+        else:
+          self._state.apply_gradients(grads)
+    self._cached_losses.clear()
+    self._accumulated_grads = None
+    self._micro_step_count = 0
+    if (
+        hasattr(self, "_learning_rate_schedule")
+        and self._learning_rate_schedule is not None
+    ):
+      try:
+        lr = self._learning_rate_schedule(self.train_step)
+        self.record_metrics("lr", lr)
+      except Exception:  # pylint: disable=broad-except
+        pass
+    self._train_step += 1
+
+  def eval_step(
+      self, payload: abstract_engine.TrainerPayload, **kwargs: Any
+  ) -> None:
+    """Executes an evaluation step on the given payload.
+
+    Args:
+      payload: Packed micro-batch evaluation input.
+      **kwargs: Additional keyword arguments for evaluation.
+    """
+
+  def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None:
+    """Forces asynchronous Orbax checkpoint serialization.
+
+    Args:
+      metadata: Checkpoint metadata payload from Orchestrator.
+      **kwargs: Additional checkpoint saving options.
+    """
+    step = kwargs.get("step", self.train_step)
+    # TODO(mazumdera): Also save self._accumulated_grads and _micro_step_count.
+    ckpt_saved = self._checkpoint_manager.save_checkpoint(
+        step=step,
+        model=self.model,
+        optimizer=self.optimizer,
+        custom_metadata=metadata,
+    )
+    if ckpt_saved:
+      logging.info("Checkpoint saved at step %d.", step)
+
+  def restore_checkpoint(self, **kwargs: Any) -> Any:
+    """Restores the latest Multi-Tier Checkpoint and returns its metadata.
+
+    Args:
+      **kwargs: Additional checkpoint restoration options.
+
+    Returns:
+      The metadata PyTree of the restored checkpoint.
+    """
+    step = kwargs.get("step", None)
+    restored_step, restored_metadata = (
+        self._checkpoint_manager.restore_checkpoint(
+            model=self.model,
+            optimizer=self.optimizer,
+            step=step,
+        )
+    )
+    if restored_step:
+      logging.info("Checkpoint restored from step %d.", restored_step)
+      self.train_step = restored_step
+    return restored_metadata
+
+  def record_metrics(
+      self,
+      name: str,
+      metric: abstract_engine.WeightedMetric | jax.Array | float | int,
+      aggregation_fn: Callable[[jax.Array], Any] | None = None,
+  ) -> None:
+    """Records a metric into the buffer, appending to JAX arrays.
+
+    Args:
+      name: The name of the metric.
+      metric: The metric to record.
+      aggregation_fn: The aggregation function to apply to the metric.
+    """
+    self._metrics_logger.buffer_metrics(
+        train_step=self.train_step,
+        name=name,
+        metric=metric,
+        aggregation_fn=aggregation_fn,
+    )
+
+  def get_metrics(
+      self, clear_cache: bool = True
+  ) -> abstract_engine.MetricsBuffer:
+    """Returns accumulated step metrics as an on-device MetricsBuffer.
+
+    Args:
+      clear_cache: Whether to reset cached metrics after retrieval.
+
+    Returns:
+      On-device MetricsBuffer containing WeightedMetric and scalar arrays.
+    """
+    return self._metrics_logger.get_metrics(clear_cache=clear_cache)
+
+  def prepare_weight_sync(self, **kwargs: Any) -> Any:
+    """Stages weights for transfer and returns access coordinates.
+
+    Args:
+      **kwargs: Weight staging parameters.
+
+    Returns:
+      Synchronization endpoints or coordinates for rollout actors.
+    """
+    return {}
+
+  def close(self) -> None:
+    """Closes the trainer and its associated resources."""
+    self._checkpoint_manager.close()
diff --git a/src/maxtext/training_engine/maxtext_engine_e2e_test.py b/src/maxtext/training_engine/maxtext_engine_e2e_test.py
new file mode 100644
index 0000000000..56bb9cbaad
--- /dev/null
+++ b/src/maxtext/training_engine/maxtext_engine_e2e_test.py
@@ -0,0 +1,175 @@
+# 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.
+
+"""End-to-end RL Orchestrator training loop driver and integration test."""
+
+from collections.abc import Iterator
+import dataclasses
+from typing import Any
+from unittest import mock
+
+from absl.testing import absltest
+from flax import nnx
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import maxtext_engine
+import optax
+
+
+class DummyNNXModel(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0]))
+
+
+@dataclasses.dataclass(kw_only=True)
+class DummyPayload(abstract_engine.TrainerPayload):
+  token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2)))
+  token_mask: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2)))
+
+
+class TrainingLoopRunner:
+  """Drives an end-to-end RL training loop across MaxTextTrainingEngine APIs."""
+
+  def __init__(
+      self,
+      trainer_instance: maxtext_engine.MaxTextTrainingEngine,
+      microbatches_per_minibatch: int = 2,
+      checkpoint_interval: int = 2,
+      eval_interval: int = 2,
+  ):
+    self.trainer = trainer_instance
+    self.microbatches_per_minibatch = microbatches_per_minibatch
+    self.checkpoint_interval = checkpoint_interval
+    self.eval_interval = eval_interval
+
+  def run(
+      self,
+      train_dataloader: Iterator[abstract_engine.TrainerPayload],
+      eval_dataloader: Iterator[abstract_engine.TrainerPayload],
+      num_minibatches: int,
+      dummy_compile_payload: abstract_engine.TrainerPayload | None = None,
+  ) -> list[abstract_engine.MetricsBuffer]:
+    """Executes the full RL training loop and returns step metric buffers."""
+    history: list[abstract_engine.MetricsBuffer] = []
+
+    _ = self.trainer.restore_checkpoint()
+    if dummy_compile_payload is not None:
+      self.trainer.compile(dummy_compile_payload)
+
+    for step in range(1, num_minibatches + 1):
+      for _ in range(self.microbatches_per_minibatch):
+        micro_payload = next(train_dataloader)
+        self.trainer.fwd_bwd(micro_payload)
+
+      self.trainer.update()
+
+      if step % self.checkpoint_interval == 0:
+        self.trainer.save_checkpoint(
+            metadata={"step": step, "source": "TrainingLoopRunner"}
+        )
+
+      if step % self.eval_interval == 0:
+        eval_payload = next(eval_dataloader)
+        self.trainer.eval_step(eval_payload)
+
+      step_metrics = self.trainer.get_metrics(clear_cache=True)
+      history.append(step_metrics)
+
+      _ = self.trainer.prepare_weight_sync()
+
+    self.trainer.close()
+    return history
+
+
+class MaxTextTrainingEngineE2ETest(absltest.TestCase):
+
+  def setup_config(self, enable_checkpointing: bool = False):
+    mock_config = mock.MagicMock(spec=pyconfig.HyperParameters)
+    mock_config.init_weights_seed = 42
+    mock_config.model_name = "llama3.1-8b"
+    mock_config.tensorboard_dir = "/tmp/tb_dir"
+    mock_config.run_name = "test_run"
+    mock_config.enable_tensorboard = False
+    if enable_checkpointing:
+      mock_config.checkpoint_directory = "/tmp/test_out/e2e_checkpoints"
+      mock_config.checkpoint_period = 2
+      mock_config.max_num_checkpoints_to_keep = 5
+      mock_config.async_checkpointing = True
+    return mock_config
+
+  @mock.patch.object(maxtext_engine.train_utils, "create_training_optimizer")
+  @mock.patch.object(maxtext_engine.checkpointing, "CheckpointManager")
+  @mock.patch.object(
+      maxtext_engine.gradient_accumulation,
+      "gradient_accumulation_loss_and_grad",
+  )
+  @mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained")
+  def test_e2e_training_loop_exercises_all_trainer_apis(
+      self, mock_from_pretrained, mock_ga, unused_mock_ckpt_mgr, mock_create_opt
+  ):
+    mock_config = self.setup_config(enable_checkpointing=True)
+    dummy_model = DummyNNXModel()
+    mock_from_pretrained.return_value = dummy_model
+    dummy_opt = nnx.Optimizer(dummy_model, optax.sgd(0.01), wrt=nnx.Param)
+    mock_create_opt.return_value = (lambda step: jnp.array(0.001), dummy_opt)
+    mock_ga.return_value = (
+        jnp.array(0.25),
+        {},
+        {"weights": jnp.array([0.1, 0.1])},
+    )
+
+    trainer_instance = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    trainer_instance._checkpoint_manager.restore_checkpoint.return_value = (
+        None,
+        {},
+    )
+    trainer_instance.with_loss_fn(
+        lambda *args, **kwargs: (
+            abstract_engine.WeightedMetric(
+                unreduced_sum=jnp.array(1.0), denominator=jnp.array(4.0)
+            ),
+            {},
+        )
+    )
+
+    runner = TrainingLoopRunner(
+        trainer_instance=trainer_instance,
+        microbatches_per_minibatch=2,
+        checkpoint_interval=2,
+        eval_interval=2,
+    )
+
+    def payload_generator() -> Iterator[abstract_engine.TrainerPayload]:
+      while True:
+        yield DummyPayload()
+
+    history = runner.run(
+        train_dataloader=payload_generator(),
+        eval_dataloader=payload_generator(),
+        num_minibatches=4,
+        dummy_compile_payload=DummyPayload(),
+    )
+
+    self.assertLen(history, 4)
+    for metrics_buf_list in history:
+      self.assertIsInstance(metrics_buf_list, list)
+      self.assertNotEmpty(metrics_buf_list)
+      self.assertIsInstance(metrics_buf_list[0], abstract_engine.MetricsBuffer)
+    self.assertEqual(trainer_instance.train_step, 4)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/src/maxtext/training_engine/maxtext_engine_test.py b/src/maxtext/training_engine/maxtext_engine_test.py
new file mode 100644
index 0000000000..771f52588f
--- /dev/null
+++ b/src/maxtext/training_engine/maxtext_engine_test.py
@@ -0,0 +1,297 @@
+# 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.
+
+import dataclasses
+from typing import Any
+from unittest import mock
+
+from absl.testing import absltest
+from flax import nnx
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import maxtext_engine
+import numpy as np
+import optax
+import orbax.checkpoint as ocp
+
+
+class DummyNNXModel(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0]))
+
+
+@dataclasses.dataclass(kw_only=True)
+class DummyPayload(abstract_engine.TrainerPayload):
+  token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2)))
+  token_mask: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2)))
+
+
+class MaxTextTrainingEngineTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    dummy_model = DummyNNXModel()
+    dummy_opt = nnx.Optimizer(dummy_model, optax.sgd(0.01), wrt=nnx.Param)
+    patcher = mock.patch.object(
+        maxtext_engine.train_utils,
+        "create_training_optimizer",
+        return_value=(lambda step: jnp.array(0.001), dummy_opt),
+    )
+    self.addCleanup(patcher.stop)
+    patcher.start()
+
+    from_pretrained_patcher = mock.patch.object(
+        maxtext_engine.model_creation_utils,
+        "from_pretrained",
+        return_value=dummy_model,
+    )
+    self.addCleanup(from_pretrained_patcher.stop)
+    self.mock_from_pretrained = from_pretrained_patcher.start()
+    self.mock_config = self.setup_config()
+
+  def setup_config(self, enable_checkpointing: bool = False):
+    mock_config = mock.MagicMock(spec=pyconfig.HyperParameters)
+    mock_config.init_weights_seed = 42
+    mock_config.model_name = "llama3.1-8b"
+    mock_config.learning_rate_schedule = mock.MagicMock(return_value=0.001)
+    mock_config.tensorboard_dir = "/tmp/tb_dir"
+    mock_config.run_name = "test_run"
+    mock_config.enable_tensorboard = False
+    if enable_checkpointing:
+      mock_config.checkpoint_directory = "/tmp/test_out/checkpoints"
+      mock_config.checkpoint_period = 500
+      mock_config.max_num_checkpoints_to_keep = 10
+      mock_config.async_checkpointing = True
+    return mock_config
+
+  def test_raises_type_error_for_non_pyconfig(self):
+    invalid_config = abstract_engine.TrainingConfig()
+    with self.assertRaises(TypeError):
+      maxtext_engine.MaxTextTrainingEngine(invalid_config)  # pytype: disable=wrong-arg-types
+
+  def test_raises_value_error_for_missing_model_name(self):
+    mock_config = mock.MagicMock(spec=pyconfig.HyperParameters)
+    mock_config.model_name = None
+    with self.assertRaises(ValueError):
+      maxtext_engine.MaxTextTrainingEngine(mock_config)
+
+  @mock.patch.object(maxtext_engine.checkpointing, "CheckpointManager")
+  @mock.patch.object(
+      maxtext_engine.gradient_accumulation,
+      "gradient_accumulation_loss_and_grad",
+  )
+  def test_max_text_trainer_instantiation_with_pyconfig(
+      self, mock_ga, unused_mock_ckpt_mgr
+  ):
+    mock_ga.return_value = (
+        jnp.array(0.5),
+        {},
+        {"weights": jnp.array([0.1, 0.2])},
+    )
+
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    self.assertIsInstance(t, abstract_engine.AbstractTrainingEngine)
+    self.mock_from_pretrained.assert_called_once()
+    self.assertEqual(t.train_step, 0)
+    payload = DummyPayload(
+        token_ids=jnp.ones((2, 2)),
+        token_mask=jnp.ones((2, 2)),
+    )
+    t.compile(payload)
+    self.assertTrue(t._compiled)
+    t.with_loss_fn(lambda *args, **kwargs: (jnp.array(0.5), {}))
+    self.assertFalse(t._compiled)
+    t.fwd_bwd(payload)
+    self.assertEqual(t._micro_step_count, 1)
+    t.update()
+    self.assertEqual(t._micro_step_count, 0)
+    self.assertIsNone(t._accumulated_grads)
+    metrics = t.get_metrics()
+    self.assertIsInstance(metrics, list)
+
+  @mock.patch("orbax.checkpoint.CheckpointManager")
+  def test_max_text_trainer_checkpoint_manager_init(self, mock_create_mgr):
+    mock_config = self.setup_config(enable_checkpointing=True)
+
+    _ = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    mock_create_mgr.assert_called_once_with(
+        directory=mock_config.checkpoint_directory,
+        options=ocp.CheckpointManagerOptions(
+            save_interval_steps=mock_config.checkpoint_period,
+            max_to_keep=mock_config.max_num_checkpoints_to_keep,
+            enable_async_checkpointing=mock_config.async_checkpointing,
+        ),
+    )
+
+  def test_save_checkpoint_with_model_and_optimizer(self):
+    mock_config = self.setup_config(enable_checkpointing=True)
+
+    t = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    mock_orbax_mgr = mock.MagicMock()
+    mock_orbax_mgr.latest_step.return_value = None
+    mock_orbax_mgr.save.return_value = True
+    t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr
+
+    t.train_step = 10
+    metadata_to_save = {"worker_id": "worker_0", "run_id": "run_123"}
+    t.save_checkpoint(metadata=metadata_to_save)
+
+    mock_orbax_mgr.save.assert_called_once()
+
+  def test_save_checkpoint_skips_if_already_saved(self):
+    mock_config = self.setup_config(enable_checkpointing=True)
+
+    t = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    mock_orbax_mgr = mock.MagicMock()
+    mock_orbax_mgr.latest_step.return_value = 10
+    t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr
+    t.train_step = 10
+
+    t.save_checkpoint(metadata={"key": "val"})
+    mock_orbax_mgr.save.assert_not_called()
+
+  def test_restore_checkpoint_no_checkpoint_returns_defaults(self):
+    mock_config = self.setup_config(enable_checkpointing=True)
+
+    t = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    mock_orbax_mgr = mock.MagicMock()
+    mock_orbax_mgr.latest_step.return_value = None
+    t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr
+
+    restored_metadata = t.restore_checkpoint()
+    self.assertEqual(restored_metadata, {})
+
+  def test_restore_checkpoint_restores_ckpt_metadata(self):
+    mock_config = self.setup_config(enable_checkpointing=True)
+
+    t = maxtext_engine.MaxTextTrainingEngine(mock_config)
+    mock_orbax_mgr = mock.MagicMock()
+    mock_orbax_mgr.latest_step.return_value = 10
+
+    dummy_metadata = mock.MagicMock()
+    mock_orbax_mgr.metadata.return_value = dummy_metadata
+    mock_orbax_mgr.restore.return_value = {
+        "model_params": {"weights": jnp.array([1.0, 2.0])}
+    }
+
+    t._checkpoint_manager._checkpoint_manager = mock_orbax_mgr
+
+    restored_metadata = t.restore_checkpoint(step=10)
+    self.assertEqual(t.train_step, 10)
+    self.assertEqual(restored_metadata, dummy_metadata)
+    mock_orbax_mgr.restore.assert_called_once()
+
+  def test_record_metrics(self):
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    # Record WeightedMetric
+    t.record_metrics(
+        name="loss",
+        metric=abstract_engine.WeightedMetric(
+            unreduced_sum=jnp.array(20.0), denominator=jnp.array(4.0)
+        ),
+    )
+    t.record_metrics(
+        name="loss",
+        metric=abstract_engine.WeightedMetric(
+            unreduced_sum=jnp.array(30.0), denominator=jnp.array(6.0)
+        ),
+    )
+
+    # Record scalar
+    t.record_metrics(
+        name="lr",
+        metric=0.002,
+        aggregation_fn=lambda x: np.round(np.asarray(x), 4),
+    )
+
+    metrics_buffer = t.get_metrics(clear_cache=True)
+    self.assertLen(metrics_buffer, 1)
+    self.assertIn("loss", metrics_buffer[0].weighted_metrics)
+    np.testing.assert_array_equal(
+        metrics_buffer[0].weighted_metrics["loss"].unreduced_sum,
+        jnp.array([20.0, 30.0]),
+    )
+    np.testing.assert_array_equal(
+        metrics_buffer[0].weighted_metrics["loss"].denominator,
+        jnp.array([4.0, 6.0]),
+    )
+    self.assertIn("lr", metrics_buffer[0].scalar_metrics)
+    np.testing.assert_array_equal(
+        metrics_buffer[0].scalar_metrics["lr"], jnp.array([0.002])
+    )
+    self.assertIn("lr", metrics_buffer[0].aggregation_fns)
+    self.assertEqual(
+        metrics_buffer[0].aggregation_fns["lr"](jnp.array([0.002])), 0.002
+    )
+
+  def test_get_metrics(self):
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    t._metrics_logger._metrics_buffer = [
+        abstract_engine.MetricsBuffer(
+            id=1,
+            weighted_metrics={
+                "loss": abstract_engine.WeightedMetric(
+                    unreduced_sum=jnp.array(50.0), denominator=jnp.array(10.0)
+                )
+            },
+            scalar_metrics={"lr": jnp.array(0.002)},
+            aggregation_fns={"lr": lambda x: np.round(np.asarray(x), 4)},
+        )
+    ]
+    metrics = t.get_metrics(clear_cache=True)
+    self.assertIn("loss", metrics[0].weighted_metrics)
+    np.testing.assert_array_equal(
+        metrics[0].weighted_metrics["loss"].unreduced_sum, jnp.array([50.0])
+    )
+    np.testing.assert_array_equal(
+        metrics[0].weighted_metrics["loss"].denominator, jnp.array([10.0])
+    )
+    self.assertIn("lr", metrics[0].scalar_metrics)
+    np.testing.assert_array_equal(
+        metrics[0].scalar_metrics["lr"], jnp.array([0.002])
+    )
+    self.assertIn("lr", metrics[0].aggregation_fns)
+    self.assertEqual(
+        metrics[0].aggregation_fns["lr"](jnp.array([0.002])), 0.002
+    )
+    self.assertEmpty(t._metrics_logger._metrics_buffer)
+
+  @mock.patch.object(
+      maxtext_engine.gradient_accumulation,
+      "gradient_accumulation_loss_and_grad",
+  )
+  def test_trainer_for_one_global_step(self, mock_ga):
+    mock_ga.return_value = (
+        jnp.array(0.5),
+        {},
+        {"weights": jnp.array([0.1, 0.1])},
+    )
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    self.assertEqual(t.train_step, 0)
+    for mini_b in range(3):
+      for _ in range(2):
+        t.fwd_bwd(DummyPayload())
+      t.update()
+      self.assertEqual(t.train_step, mini_b + 1)
+      t.save_checkpoint(metadata={"batch": mini_b})
+    metrics = t.get_metrics(clear_cache=True)
+    self.assertLen(metrics, 3)
+    for metric in metrics:
+      self.assertIn("lr", metric.scalar_metrics)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/src/maxtext/training_engine/metrics.py b/src/maxtext/training_engine/metrics.py
new file mode 100644
index 0000000000..f454088285
--- /dev/null
+++ b/src/maxtext/training_engine/metrics.py
@@ -0,0 +1,204 @@
+# 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.
+
+"""Metric utilities for MaxRL."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+import dataclasses
+from typing import Any
+
+from absl import logging
+import jax
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.utils import max_utils
+import numpy as np
+
+
+class MetricsLogger:
+  """Logger and buffering service for metrics."""
+
+  def __init__(self, config: pyconfig.HyperParameters):
+    self._metrics_buffer: list[abstract_engine.MetricsBuffer] = []
+    self.tb_writer = max_utils.initialize_summary_writer(
+        config.tensorboard_dir, config.run_name, config.enable_tensorboard
+    )
+
+  def buffer_metrics(
+      self,
+      train_step: int,
+      name: str,
+      metric: jax.Array | float | int | None = None,
+      aggregation_fn: Callable[[jax.Array], Any] | None = None,
+  ):
+    """Buffers metrics for the given train step.
+
+    Args:
+      train_step: The train step for which to buffer metrics.
+      name: The name of the metric.
+      metric: The metric to buffer.
+      aggregation_fn: Optional aggregation function to apply to the metric.
+    """
+    if not self._metrics_buffer or self._metrics_buffer[-1].id != train_step:
+      new_buffer = abstract_engine.MetricsBuffer(id=train_step, mode="train")
+      self._metrics_buffer.append(new_buffer)
+
+    # Record the new metric in the buffer for the current step.
+    self._record_metric(name, metric, aggregation_fn=aggregation_fn)
+
+  def _record_metric(
+      self,
+      name: str,
+      metric: abstract_engine.WeightedMetric | jax.Array | float | int,
+      aggregation_fn: Callable[[jax.Array], Any] | None = None,
+  ) -> None:
+    """Records a metric into the buffer, appending to JAX arrays.
+
+    Args:
+      name: The name of the metric.
+      metric: The metric to record.
+      aggregation_fn: The aggregation function to apply to the metric.
+    """
+    buffer = self._metrics_buffer[-1]
+
+    if aggregation_fn is not None:
+      buffer.aggregation_fns[name] = aggregation_fn
+
+    if isinstance(metric, abstract_engine.WeightedMetric):
+      if name not in buffer.weighted_metrics:
+        buffer.weighted_metrics[name] = abstract_engine.WeightedMetric(
+            unreduced_sum=jax.numpy.atleast_1d(metric.unreduced_sum),
+            denominator=jax.numpy.atleast_1d(metric.denominator),
+            eps=metric.eps,
+            min_denom=metric.min_denom,
+        )
+      else:
+        current_val = buffer.weighted_metrics[name]
+        new_sum = jax.numpy.append(
+            current_val.unreduced_sum, metric.unreduced_sum
+        )
+        new_denom = jax.numpy.append(
+            current_val.denominator, metric.denominator
+        )
+        buffer.weighted_metrics[name] = dataclasses.replace(
+            current_val, unreduced_sum=new_sum, denominator=new_denom
+        )
+    else:
+      if name not in buffer.scalar_metrics:
+        buffer.scalar_metrics[name] = jax.numpy.atleast_1d(
+            jax.numpy.asarray(metric)
+        )
+      else:
+        buffer.scalar_metrics[name] = jax.numpy.append(
+            buffer.scalar_metrics[name], metric
+        )
+    self._metrics_buffer[-1] = buffer
+
+  def get_metrics(
+      self, clear_cache: bool = True
+  ) -> list[abstract_engine.MetricsBuffer]:
+    """Returns cached metrics and optionally clears the metrics cache.
+
+    Args:
+      clear_cache: Whether to reset cached metrics after retrieval.
+
+    Returns:
+      The accumulated on-device MetricsBuffer.
+    """
+    metrics_to_return = self._metrics_buffer
+    if clear_cache:
+      # Clear the metrics buffer.
+      self._metrics_buffer = []
+    return metrics_to_return
+
+  def log_metrics(self, metrics: dict[str, Any]) -> None:
+    """Log metrics.
+
+    Args:
+      metrics: Dictionary mapping metric names to reduced Python floats/numpy
+        arrays.
+    """
+    for metric_name, value in metrics.items():
+      try:
+        agg_value = np.array(value)
+        logging.info("Metric %s: %s", metric_name, agg_value)
+      except Exception:  # pylint: disable=broad-except
+        logging.warning(
+            "Skipping metric %s: Could not convert to numpy array.", metric_name
+        )
+        continue
+
+  def write_metrics(self, metrics: dict[str, Any]) -> None:
+    """Write metrics to TensorBoard.
+
+    Args:
+      metrics: Dictionary mapping metric names to reduced Python floats/numpy
+        arrays.
+    """
+    self.log_metrics(metrics)
+    self.write_metrics_to_tensorboard(metrics)
+
+  def write_metrics_to_tensorboard(self, metrics: dict[str, Any]) -> None:
+    """Write metrics to TensorBoard.
+
+    Args:
+      metrics: Dictionary mapping metric names to reduced Python floats/numpy
+        arrays.
+    """
+    if jax.process_index() == 0:
+      for metric_name, value in metrics.items():
+        self.tb_writer.add_scalar(metric_name, value)
+        self.tb_writer.flush()
+
+  def process_metrics(self, metrics: abstract_engine.MetricsBuffer) -> None:
+    """Reduction and processing pipeline for MetricsBuffer.
+
+    Unpacks on-device weighted metrics by invoking safe compute(), extracts
+    scalar metrics, and applies any host-side aggregation callbacks.
+
+    Args:
+      metrics: MetricsBuffer retrieved from device.
+    """
+    processed: dict[str, Any] = {}
+
+    # Process weighted metrics via safe division
+    for name, weighted_metric in metrics.weighted_metrics.items():
+      reduced_val = weighted_metric.compute()
+      host_val = np.asarray(reduced_val)
+      if host_val.ndim == 0:
+        host_val = float(host_val)
+      if name in metrics.aggregation_fns:
+        host_val = metrics.aggregation_fns[name](host_val)
+      processed[name] = host_val
+
+    # Process scalar metrics
+    for name, scalar_val in metrics.scalar_metrics.items():
+      host_val = np.asarray(scalar_val)
+      if host_val.ndim == 0:
+        host_val = float(host_val)
+      if name in metrics.aggregation_fns:
+        host_val = metrics.aggregation_fns[name](host_val)
+      processed[name] = host_val
+
+    logging.info("Logging metrics for step %s", metrics.id)
+    self.write_metrics(processed)
+
+  def flush_metrics_and_cleanup(self) -> list[abstract_engine.MetricsBuffer]:
+    """Flushes metrics and cleans up the metrics buffer."""
+    if self._metrics_buffer is None:
+      return None
+    self.process_metrics(self._metrics_buffer)
+    self._metrics_buffer = None