rlx_rl/graph/
actor.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 actor graphs.
16
17use rlx_compile::legalize_broadcast;
18use rlx_ir::{DType, Graph, NodeId, Op, Shape};
19
20use crate::graph::mlp::{
21    ParamSlot, concat_features, flow_map_jump, init_mat, init_vec, mlp_layers, mse_mean,
22};
23use crate::spec::RlSpec;
24
25/// Named parameter tensors shared across compiled graphs.
26#[derive(Debug, Clone, Default)]
27pub struct WeightStore(pub std::collections::HashMap<String, Vec<f32>>);
28
29impl WeightStore {
30    pub fn apply(&self, exec: &mut rlx_runtime::CompiledGraph) {
31        for (name, data) in &self.0 {
32            exec.set_param(name, data);
33        }
34    }
35}
36
37#[derive(Debug, Clone)]
38pub struct ActorGraphBundle {
39    pub forward: Graph,
40    /// Average velocity \(u_{r,t}\) only (FMQ anchor / online targets).
41    pub velocity: Graph,
42    pub offline_train: ActorTrainGraph,
43    pub online_train: ActorTrainGraph,
44}
45
46#[derive(Debug, Clone)]
47pub struct ActorTrainGraph {
48    pub graph: Graph,
49    pub d_output: NodeId,
50    pub loss: NodeId,
51    pub params: Vec<ParamSlot>,
52}
53
54pub fn build_actor_graphs(spec: &RlSpec) -> ActorGraphBundle {
55    ActorGraphBundle {
56        forward: legalize_broadcast::run(build_forward(spec)),
57        velocity: legalize_broadcast::run(build_velocity(spec)),
58        offline_train: build_offline_train(spec),
59        online_train: build_online_train(spec),
60    }
61}
62
63pub struct CompiledFlowMapAgent {
64    pub forward: rlx_runtime::CompiledGraph,
65    pub velocity: rlx_runtime::CompiledGraph,
66    pub offline: rlx_runtime::CompiledGraph,
67    pub online: rlx_runtime::CompiledGraph,
68    pub offline_meta: ActorTrainGraph,
69    pub online_meta: ActorTrainGraph,
70}
71
72impl CompiledFlowMapAgent {
73    pub fn compile(session: &rlx_runtime::Session, bundle: &ActorGraphBundle) -> Self {
74        Self {
75            forward: session.compile(bundle.forward.clone()),
76            velocity: session.compile(bundle.velocity.clone()),
77            offline: session.compile(bundle.offline_train.graph.clone()),
78            online: session.compile(bundle.online_train.graph.clone()),
79            offline_meta: bundle.offline_train.clone(),
80            online_meta: bundle.online_train.clone(),
81        }
82    }
83
84    pub fn set_weights(&mut self, weights: &WeightStore) {
85        for g in [
86            &mut self.forward,
87            &mut self.velocity,
88            &mut self.offline,
89            &mut self.online,
90        ] {
91            weights.apply(g);
92        }
93    }
94
95    /// Average velocity \(u_{r,t}(a_r|s)\) (Python `actor_bc_flow` output).
96    pub fn velocity(&mut self, state: &[f32], a_r: &[f32], r: f32, t: f32) -> Vec<f32> {
97        self.velocity
98            .run(&[("state", state), ("a_r", a_r), ("r", &[r]), ("t", &[t])])
99            .into_iter()
100            .next()
101            .unwrap_or_default()
102    }
103
104    pub fn one_step(&mut self, state: &[f32], a0: &[f32]) -> Vec<f32> {
105        let r = [0.0f32];
106        let t = [1.0f32];
107        self.forward
108            .run(&[("state", state), ("a_r", a0), ("r", &r), ("t", &t)])
109            .into_iter()
110            .next()
111            .unwrap_or_default()
112    }
113
114    pub fn jump(&mut self, state: &[f32], a_r: &[f32], r: f32, t: f32) -> Vec<f32> {
115        let rv = [r];
116        let tv = [t];
117        self.forward
118            .run(&[("state", state), ("a_r", a_r), ("r", &rv), ("t", &tv)])
119            .into_iter()
120            .next()
121            .unwrap_or_default()
122    }
123}
124
125pub fn init_actor_weights(spec: &RlSpec, seed: u64) -> WeightStore {
126    let mut w = WeightStore::default();
127    let mut s = seed;
128    let mut in_d = spec.actor_in_dim();
129    for (li, &hd) in spec.hidden.iter().enumerate() {
130        init_mat(&mut w, &format!("actor_w{li}"), in_d, hd, &mut s);
131        init_vec(&mut w, &format!("actor_b{li}"), hd, &mut s);
132        in_d = hd;
133    }
134    init_mat(&mut w, "actor_w_out", in_d, spec.action_dim, &mut s);
135    init_vec(&mut w, "actor_b_out", spec.action_dim, &mut s);
136    w
137}
138
139fn build_forward(spec: &RlSpec) -> Graph {
140    let f = DType::F32;
141    let b = spec.batch;
142    let sd = spec.state_dim;
143    let ad = spec.action_dim;
144    let mut g = Graph::new("flow_map_forward");
145    let mut params = Vec::new();
146
147    let state = g.input("state", Shape::new(&[b, sd], f));
148    let a_r = g.input("a_r", Shape::new(&[b, ad], f));
149    let r = g.input("r", Shape::new(&[b, 1], f));
150    let t = g.input("t", Shape::new(&[b, 1], f));
151
152    let feats = concat_features(&mut g, vec![state, a_r, r, t]);
153    let u = mlp_layers(
154        &mut g,
155        feats,
156        spec.actor_in_dim(),
157        &spec.hidden,
158        ad,
159        "actor",
160        &mut params,
161    );
162    let a_t = flow_map_jump(&mut g, a_r, u, r, t, b);
163    let _ = params;
164    g.set_outputs(vec![a_t]);
165    g
166}
167
168fn build_velocity(spec: &RlSpec) -> Graph {
169    let f = DType::F32;
170    let b = spec.batch;
171    let sd = spec.state_dim;
172    let ad = spec.action_dim;
173    let mut g = Graph::new("flow_map_velocity");
174    let mut params = Vec::new();
175
176    let state = g.input("state", Shape::new(&[b, sd], f));
177    let a_r = g.input("a_r", Shape::new(&[b, ad], f));
178    let r = g.input("r", Shape::new(&[b, 1], f));
179    let t = g.input("t", Shape::new(&[b, 1], f));
180
181    let feats = concat_features(&mut g, vec![state, a_r, r, t]);
182    let u = mlp_layers(
183        &mut g,
184        feats,
185        spec.actor_in_dim(),
186        &spec.hidden,
187        ad,
188        "actor",
189        &mut params,
190    );
191    let _ = params;
192    g.set_outputs(vec![u]);
193    g
194}
195
196fn build_offline_train(spec: &RlSpec) -> ActorTrainGraph {
197    let f = DType::F32;
198    let b = spec.batch;
199    let sd = spec.state_dim;
200    let ad = spec.action_dim;
201    let mut g = Graph::new("flow_map_offline");
202    let mut params = Vec::new();
203
204    let state = g.input("state", Shape::new(&[b, sd], f));
205    let a_r = g.input("a_r", Shape::new(&[b, ad], f));
206    let r = g.input("r", Shape::new(&[b, 1], f));
207    let t = g.input("t", Shape::new(&[b, 1], f));
208    let target_u = g.input("target_u", Shape::new(&[b, ad], f));
209
210    let feats = concat_features(&mut g, vec![state, a_r, r, t]);
211    let u = mlp_layers(
212        &mut g,
213        feats,
214        spec.actor_in_dim(),
215        &spec.hidden,
216        ad,
217        "actor",
218        &mut params,
219    );
220    let loss = mse_mean(&mut g, u, target_u);
221    g.set_outputs(vec![loss]);
222    finalize_train(g, params, loss)
223}
224
225fn build_online_train(spec: &RlSpec) -> ActorTrainGraph {
226    let f = DType::F32;
227    let b = spec.batch;
228    let sd = spec.state_dim;
229    let ad = spec.action_dim;
230    let mut g = Graph::new("flow_map_online");
231    let mut params = Vec::new();
232
233    let state = g.input("state", Shape::new(&[b, sd], f));
234    let a_r = g.input("a_r", Shape::new(&[b, ad], f));
235    let target_u = g.input("target_u", Shape::new(&[b, ad], f));
236    let r = g.input("r", Shape::new(&[b, 1], f));
237    let t = g.input("t", Shape::new(&[b, 1], f));
238
239    let feats = concat_features(&mut g, vec![state, a_r, r, t]);
240    let u = mlp_layers(
241        &mut g,
242        feats,
243        spec.actor_in_dim(),
244        &spec.hidden,
245        ad,
246        "actor",
247        &mut params,
248    );
249    let loss = mse_mean(&mut g, u, target_u);
250    g.set_outputs(vec![loss]);
251    finalize_train(g, params, loss)
252}
253
254fn finalize_train(g: Graph, params: Vec<ParamSlot>, loss_fwd: NodeId) -> ActorTrainGraph {
255    let (g, remap) = legalize_broadcast::run_with_remap(g);
256    let params: Vec<ParamSlot> = params
257        .into_iter()
258        .map(|mut p| {
259            p.param = remap[&p.param];
260            p
261        })
262        .collect();
263    let loss_fwd = remap[&loss_fwd];
264    let wrt: Vec<NodeId> = params.iter().map(|p| p.param).collect();
265    let bwd = rlx_autodiff::grad_with_loss(&g, &wrt);
266    let d_output = bwd
267        .nodes()
268        .iter()
269        .find(|n| matches!(&n.op, Op::Input { name } if name == "d_output"))
270        .map(|n| n.id)
271        .expect("d_output");
272    let grad_ids: Vec<NodeId> = bwd.outputs[1..=params.len()].to_vec();
273    let params: Vec<ParamSlot> = params
274        .into_iter()
275        .zip(grad_ids)
276        .map(|(mut p, grad)| {
277            p.grad = Some(grad);
278            p
279        })
280        .collect();
281    ActorTrainGraph {
282        graph: bwd,
283        d_output,
284        loss: loss_fwd,
285        params,
286    }
287}
rlx_rl/graph/actor.rs

rlx_rl/graph/
actor.rs
