Enhance RLConfig Pydantic inheritance hierarchy and establish unified import and gradient accumulation conventions.

src/maxtext/configs/types.py

@@ -3982,9 +3982,32 @@ class RLConfig(
     Decoding,
     IciParallelism,
     DcnParallelism,
+    PipelineParallelism,
+    DilocoParams,
     HardwareAndMesh,
     ModelArchitecture,
+    MTP,
     MoBa,
+    # Advanced Architectures, Tuning, and Optimizers
+    Muon,
+    FineTuning,
+    Distillation,
+    # Datasets and Loading Compatibility
+    DatasetGeneral,
+    TfdsDataset,
+    HfDataset,
+    GrainDataset,
+    OlmoGrainDataset,
+    # Inference, Checkpointing, and Monitoring
+    EmergencyCheckpointing,
+    ElasticTraining,
+    InferenceServer,
+    InferenceBenchmark,
+    PrefixCaching,
+    HloDump,
+    Goodput,
+    GcpMonitoring,
+    ManagedMLDiagnostics,
     # Positional Embeddings
     PositionalEmbedding,
     Rope,
@@ -4011,9 +4034,12 @@ class RLConfig(
     AttentionIndexer,
     SplashAttention,
     Qwen3Next,
-    # Debugging and Profiling
+    # Debugging, Profiling, and Telemetry
+    AOT,
     DevelopmentAndDebugging,
     Profiling,
+    Metrics,
+    Tensorboard,
     # For compatibility with trainer in post_train/rl
     RL,
     RLCluster,
@@ -4022,6 +4048,8 @@ class RLConfig(
     RLReward,
     RLSpecialTokens,
     VLLM,
+    TrainingLoop,
+    DerivedValues,
 ):
   """
   Configuration for Reinforcement Learning in MaxText.
@@ -4207,19 +4235,20 @@ def set_derived_values_and_validate(self) -> "RLConfig":
 
     # Dynamically inject model dimensions.
     emb_scale, num_head_scale, mlp_dim_scale, layer_scale = get_individual_scales(self.global_parameter_scale)
-    object.__setattr__(self, "emb_dim", int((2**emb_scale) * self.base_emb_dim))
-    object.__setattr__(self, "num_query_heads", int((2**num_head_scale) * self.base_num_query_heads))
-    object.__setattr__(self, "num_kv_heads", int((2**num_head_scale) * self.base_num_kv_heads))
-    object.__setattr__(self, "mlp_dim", int((2**mlp_dim_scale) * self.base_mlp_dim))
-    object.__setattr__(self, "moe_mlp_dim", int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0)))
-    object.__setattr__(self, "num_decoder_layers", int((2**layer_scale) * self.base_num_decoder_layers))
+    self.emb_dim = int((2**emb_scale) * self.base_emb_dim)
+    self.num_query_heads = int((2**num_head_scale) * self.base_num_query_heads)
+    self.num_kv_heads = int((2**num_head_scale) * self.base_num_kv_heads)
+    self.mlp_dim = int((2**mlp_dim_scale) * self.base_mlp_dim)
+    self.moe_mlp_dim = int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0))
+    self.num_decoder_layers = int((2**layer_scale) * self.base_num_decoder_layers)
 
     # Mirror into internal MaxText fields for backward compatibility.
     train_micro_batch_size = getattr(self.dataset, "train_micro_batch_size", -1)
     batch_size = getattr(self.dataset, "batch_size", 1)
     if train_micro_batch_size <= 0:
       train_micro_batch_size = batch_size
-    object.__setattr__(self, "micro_batch_size_to_train_on", train_micro_batch_size)
+    self.micro_batch_size_to_train_on = train_micro_batch_size
+    self.steps = getattr(self, "train_steps", getattr(self, "num_batches", 10))
 
     if self.remat_policy == "custom":
       tensors = [
@@ -4244,7 +4273,7 @@ def set_derived_values_and_validate(self) -> "RLConfig":
           "attention_out",
           "out_proj",
       ]
-      object.__setattr__(self, "tensors_on_device", [t for t in tensors if getattr(self, t) == "device"])
-      object.__setattr__(self, "tensors_to_offload", [t for t in tensors if getattr(self, t) == "offload"])
+      self.tensors_on_device = [t for t in tensors if getattr(self, t) == "device"]
+      self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"]
 
     return self