rlx_rl/
dataset.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 — offline dataset batches.
16
17use crate::buffer::{ReplayBuffer, Transition};
18use crate::flow_curriculum::sample_r_t;
19use crate::spec::RlSpec;
20
21/// Static offline transitions for flow-matching pretraining.
22#[derive(Debug, Clone, Default)]
23pub struct OfflineDataset {
24    pub transitions: Vec<Transition>,
25}
26
27impl OfflineDataset {
28    pub fn from_transitions(transitions: Vec<Transition>) -> Self {
29        Self { transitions }
30    }
31
32    pub fn len(&self) -> usize {
33        self.transitions.len()
34    }
35
36    pub fn is_empty(&self) -> bool {
37        self.transitions.is_empty()
38    }
39
40    pub fn from_replay(buf: &ReplayBuffer) -> Self {
41        Self {
42            transitions: buf.iter().cloned().collect(),
43        }
44    }
45
46    /// L_Diag batch (Eq. 3): `r = t`, target velocity `a₁ - a₀`.
47    pub fn sample_diag_batch(
48        &self,
49        spec: &RlSpec,
50        indices: &[usize],
51        rng: &mut u64,
52    ) -> ActorTrainBatch {
53        let b = indices.len();
54        let sd = spec.state_dim;
55        let ad = spec.action_dim;
56        let mut state = vec![0.0f32; b * sd];
57        let mut a_r = vec![0.0f32; b * ad];
58        let mut r = vec![0.0f32; b];
59        let mut t = vec![0.0f32; b];
60        let mut target_u = vec![0.0f32; b * ad];
61
62        for (bi, &idx) in indices.iter().enumerate() {
63            let tr = &self.transitions[idx % self.transitions.len()];
64            state[bi * sd..(bi + 1) * sd].copy_from_slice(&tr.state);
65            let mut x0 = vec![0.0f32; ad];
66            for d in 0..ad {
67                x0[d] = box_muller(rng);
68            }
69            let tc = uniform01(rng);
70            r[bi] = tc;
71            t[bi] = tc;
72            for d in 0..ad {
73                a_r[bi * ad + d] = (1.0 - tc) * x0[d] + tc * tr.action[d];
74                target_u[bi * ad + d] = tr.action[d] - x0[d];
75            }
76        }
77        ActorTrainBatch {
78            state,
79            a_r,
80            r,
81            t,
82            target_u,
83        }
84    }
85
86    /// Pack ESD inputs: states, bridge `a_r`, curriculum `(r,t)`, noise velocity `v_rt`.
87    pub fn sample_esd_inputs(
88        &self,
89        spec: &RlSpec,
90        indices: &[usize],
91        step: usize,
92        rng: &mut u64,
93    ) -> EsdInputsBatch {
94        let b = indices.len();
95        let sd = spec.state_dim;
96        let ad = spec.action_dim;
97        let mut state = vec![0.0f32; b * sd];
98        let mut a_r = vec![0.0f32; b * ad];
99        let mut v_rt = vec![0.0f32; b * ad];
100        let mut gamma = vec![0.0f32; b];
101
102        let (rs, ts) = sample_r_t(
103            b,
104            step,
105            spec.flow_map_warmup_steps,
106            spec.flow_map_anneal_end_step,
107            rng,
108        );
109
110        for (bi, &idx) in indices.iter().enumerate() {
111            let tr = &self.transitions[idx % self.transitions.len()];
112            state[bi * sd..(bi + 1) * sd].copy_from_slice(&tr.state);
113            let mut x0 = vec![0.0f32; ad];
114            for d in 0..ad {
115                x0[d] = box_muller(rng);
116            }
117            let ri = rs[bi];
118            for d in 0..ad {
119                a_r[bi * ad + d] = (1.0 - ri) * x0[d] + ri * tr.action[d];
120                v_rt[bi * ad + d] = tr.action[d] - x0[d];
121            }
122            gamma[bi] = uniform01(rng);
123        }
124
125        EsdInputsBatch {
126            state,
127            a_r,
128            r: rs,
129            t: ts,
130            v_rt,
131            gamma,
132        }
133    }
134
135    /// Legacy CFM batch (diagonal only).
136    pub fn sample_cfm_batch(&self, spec: &RlSpec, indices: &[usize], rng: &mut u64) -> CfmBatch {
137        let batch = self.sample_diag_batch(spec, indices, rng);
138        CfmBatch {
139            state: batch.state,
140            a0: vec![],
141            a1: vec![],
142            a_r: batch.a_r,
143            r: batch.r,
144            t: batch.t,
145            target_u: batch.target_u,
146        }
147    }
148}
149
150#[derive(Debug, Clone)]
151pub struct ActorTrainBatch {
152    pub state: Vec<f32>,
153    pub a_r: Vec<f32>,
154    pub r: Vec<f32>,
155    pub t: Vec<f32>,
156    pub target_u: Vec<f32>,
157}
158
159#[derive(Debug, Clone)]
160pub struct EsdInputsBatch {
161    pub state: Vec<f32>,
162    pub a_r: Vec<f32>,
163    pub r: Vec<f32>,
164    pub t: Vec<f32>,
165    pub v_rt: Vec<f32>,
166    pub gamma: Vec<f32>,
167}
168
169#[derive(Debug, Clone)]
170pub struct CfmBatch {
171    pub state: Vec<f32>,
172    pub a0: Vec<f32>,
173    pub a1: Vec<f32>,
174    pub a_r: Vec<f32>,
175    pub r: Vec<f32>,
176    pub t: Vec<f32>,
177    pub target_u: Vec<f32>,
178}
179
180fn uniform01(seed: &mut u64) -> f32 {
181    *seed = crate::buffer::rand_like(*seed);
182    (*seed >> 11) as f32 / (1u32 << 21) as f32
183}
184
185fn box_muller(seed: &mut u64) -> f32 {
186    let u1 = uniform01(seed).max(1e-7);
187    let u2 = uniform01(seed);
188    (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
189}
rlx_rl/dataset.rs

rlx_rl/
dataset.rs
