rlx_rl/
spec.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15// RLX — Flow Map Q-Guidance.
16
17/// Off-diagonal self-distillation variant (Python `distillation_type`).
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
19pub enum DistillationType {
20    #[default]
21    Mf,
22    Lsd,
23    Psd,
24}
25
26impl DistillationType {
27    pub fn from_str(s: &str) -> Self {
28        match s {
29            "lsd" => Self::Lsd,
30            "psd" => Self::Psd,
31            _ => Self::Mf,
32        }
33    }
34}
35
36/// Problem dimensions and training hyperparameters.
37#[derive(Debug, Clone)]
38pub struct RlSpec {
39    pub state_dim: usize,
40    pub action_dim: usize,
41    pub batch: usize,
42    /// Hidden layers for actor and critic MLPs (excluding output heads).
43    pub hidden: Vec<usize>,
44    pub gamma: f32,
45    /// Trust-region radius for FMQ projection (before adaptive scaling).
46    pub eta: f32,
47    /// Adaptive trust region: `eta_eff = 1 / (1 + beta * delta_norm)`.
48    pub eta_beta: f32,
49    pub eta_kappa: f32,
50    pub actor_lr: f32,
51    pub critic_lr: f32,
52    pub tau: f32,
53    /// FMQ trust region: η = σ²/(2α) when `fmq_eta_override` is None (Python `fmq.py`).
54    pub fmq_alpha: f32,
55    pub fmq_sigma_sq: f32,
56    /// If set, fixes η (Python `fmq_eta_override >= 0`).
57    pub fmq_eta_override: Option<f32>,
58    pub fmq_adaptive_eta: bool,
59    /// β in Eq. 13 when `fmq_adaptive_eta` (Python `fmq_beta`).
60    pub fmq_beta: f32,
61    /// Evaluate ∇Q at `a1_online` instead of `a1_off` (Python `fmq_grad_at_online`).
62    pub fmq_grad_at_online: bool,
63    pub fmq_normalize_grad: bool,
64    /// Best-of-N sample count at eval (Python `actor_num_samples`).
65    pub actor_num_samples: usize,
66    pub action_clip: f32,
67    /// Offline / online ESD curriculum (Python `flow_map_*` / `esd_*`).
68    pub flow_map_warmup_steps: usize,
69    pub flow_map_anneal_end_step: usize,
70    pub esd_warmup_steps: usize,
71    pub esd_anneal_end_step: usize,
72    pub distillation_type: DistillationType,
73    /// Online FMQ auxiliary weights (Python `esd_weight`, `diag_weight`).
74    pub esd_weight: f32,
75    pub diag_weight: f32,
76    /// QGBS trust region at eval (Python `qgbs_eta`).
77    pub qgbs_eta: f32,
78}
79
80impl RlSpec {
81    pub fn toy(batch: usize) -> Self {
82        Self {
83            state_dim: 4,
84            action_dim: 2,
85            batch,
86            hidden: vec![64, 64],
87            gamma: 0.99,
88            eta: 0.5,
89            eta_beta: 0.3,
90            eta_kappa: 1e-4,
91            actor_lr: 3e-4,
92            critic_lr: 3e-4,
93            tau: 0.005,
94            fmq_alpha: 1.0,
95            fmq_sigma_sq: 1.0,
96            fmq_eta_override: None,
97            fmq_adaptive_eta: false,
98            fmq_beta: 0.3,
99            fmq_grad_at_online: false,
100            fmq_normalize_grad: true,
101            actor_num_samples: 32,
102            action_clip: 1.0,
103            flow_map_warmup_steps: 5,
104            flow_map_anneal_end_step: 50,
105            esd_warmup_steps: 0,
106            esd_anneal_end_step: 50,
107            distillation_type: DistillationType::Mf,
108            esd_weight: 0.0,
109            diag_weight: 0.0,
110            qgbs_eta: 0.3,
111        }
112    }
113
114    /// Trust-region step size (Python `_get_eta`).
115    pub fn fmq_eta(&self) -> f32 {
116        if let Some(e) = self.fmq_eta_override {
117            return e;
118        }
119        self.fmq_sigma_sq / (2.0 * self.fmq_alpha)
120    }
121
122    pub fn with_batch(&self, batch: usize) -> Self {
123        let mut s = self.clone();
124        s.batch = batch;
125        s
126    }
127
128    pub fn actor_in_dim(&self) -> usize {
129        self.state_dim + self.action_dim + 2
130    }
131
132    pub fn critic_in_dim(&self) -> usize {
133        self.state_dim + self.action_dim
134    }
135}
rlx_rl/spec.rs

rlx_rl/
spec.rs
