rlx_rl/graph/
critic.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 — twin critic graphs.
16
17use rlx_compile::legalize_broadcast;
18use rlx_ir::infer::GraphExt;
19use rlx_ir::{DType, Graph, NodeId, Op, Shape};
20
21use crate::graph::actor::WeightStore;
22use crate::graph::mlp::{ParamSlot, concat_features, init_mat, init_vec, mlp_trunk};
23use crate::spec::RlSpec;
24
25#[derive(Debug, Clone)]
26pub struct CriticGraphBundle {
27    pub forward: Graph,
28    pub train: CriticTrainGraph,
29    pub q_grad: CriticQGradGraph,
30}
31
32#[derive(Debug, Clone)]
33pub struct CriticTrainGraph {
34    pub graph: Graph,
35    pub d_output: NodeId,
36    pub loss: NodeId,
37    pub params: Vec<ParamSlot>,
38}
39
40#[derive(Debug, Clone)]
41pub struct CriticQGradGraph {
42    pub graph: Graph,
43    pub d_output: NodeId,
44    pub action: NodeId,
45    pub grad_action: NodeId,
46}
47
48pub fn build_critic_graphs(spec: &RlSpec) -> CriticGraphBundle {
49    CriticGraphBundle {
50        forward: legalize_broadcast::run(build_forward(spec)),
51        train: build_train(spec),
52        q_grad: build_q_grad(spec),
53    }
54}
55
56pub struct CompiledTwinCritic {
57    pub forward: rlx_runtime::CompiledGraph,
58    pub train: rlx_runtime::CompiledGraph,
59    pub q_grad: rlx_runtime::CompiledGraph,
60    pub train_meta: CriticTrainGraph,
61    pub q_grad_meta: CriticQGradGraph,
62}
63
64impl CompiledTwinCritic {
65    pub fn compile(session: &rlx_runtime::Session, bundle: &CriticGraphBundle) -> Self {
66        Self {
67            forward: session.compile(bundle.forward.clone()),
68            train: session.compile(bundle.train.graph.clone()),
69            q_grad: session.compile(bundle.q_grad.graph.clone()),
70            train_meta: bundle.train.clone(),
71            q_grad_meta: bundle.q_grad.clone(),
72        }
73    }
74
75    pub fn set_weights(&mut self, weights: &WeightStore) {
76        for g in [&mut self.forward, &mut self.train, &mut self.q_grad] {
77            weights.apply(g);
78        }
79    }
80
81    pub fn q_values(&mut self, state: &[f32], action: &[f32]) -> (f32, f32) {
82        let outs = self.forward.run(&[("state", state), ("action", action)]);
83        let q1 = outs.first().map(|v| v[0]).unwrap_or(0.0);
84        let q2 = outs.get(1).map(|v| v[0]).unwrap_or(0.0);
85        (q1, q2)
86    }
87
88    pub fn action_grad(&mut self, state: &[f32], action: &[f32]) -> Vec<f32> {
89        let outs = self.q_grad.run(&[
90            ("state", state),
91            ("action", action),
92            ("d_output", &[1.0f32]),
93        ]);
94        outs.into_iter().nth(1).unwrap_or_default()
95    }
96}
97
98pub fn init_critic_weights(spec: &RlSpec, seed: u64) -> WeightStore {
99    let mut w = WeightStore::default();
100    let mut s = seed.wrapping_add(0x000C_171C);
101    let mut in_d = spec.critic_in_dim();
102    for (li, &hd) in spec.hidden.iter().enumerate() {
103        init_mat(&mut w, &format!("critic_w{li}"), in_d, hd, &mut s);
104        init_vec(&mut w, &format!("critic_b{li}"), hd, &mut s);
105        in_d = hd;
106    }
107    init_mat(&mut w, "critic_q1_w", in_d, 1, &mut s);
108    init_vec(&mut w, "critic_q1_b", 1, &mut s);
109    init_mat(&mut w, "critic_q2_w", in_d, 1, &mut s);
110    init_vec(&mut w, "critic_q2_b", 1, &mut s);
111    w
112}
113
114fn build_forward(spec: &RlSpec) -> Graph {
115    let f = DType::F32;
116    let b = spec.batch;
117    let sd = spec.state_dim;
118    let ad = spec.action_dim;
119    let mut g = Graph::new("critic_forward");
120    let mut params = Vec::new();
121
122    let state = g.input("state", Shape::new(&[b, sd], f));
123    let action = g.input("action", Shape::new(&[b, ad], f));
124    let feats = concat_features(&mut g, vec![state, action]);
125    let (h, in_d) = mlp_trunk(
126        &mut g,
127        feats,
128        spec.critic_in_dim(),
129        &spec.hidden,
130        "critic",
131        &mut params,
132    );
133    let q1 = q_head(&mut g, h, in_d, "critic_q1", b, &mut params);
134    let q2 = q_head(&mut g, h, in_d, "critic_q2", b, &mut params);
135    let _ = params;
136    g.set_outputs(vec![q1, q2]);
137    g
138}
139
140fn q_head(
141    g: &mut Graph,
142    h: NodeId,
143    in_d: usize,
144    prefix: &str,
145    batch: usize,
146    params: &mut Vec<ParamSlot>,
147) -> NodeId {
148    let f = DType::F32;
149    let w_name = format!("{prefix}_w");
150    let b_name = format!("{prefix}_b");
151    let w = g.param(&w_name, Shape::new(&[in_d, 1], f));
152    let b = g.param(&b_name, Shape::new(&[1], f));
153    params.push(ParamSlot {
154        name: w_name,
155        shape: vec![in_d, 1],
156        param: w,
157        grad: None,
158    });
159    params.push(ParamSlot {
160        name: b_name,
161        shape: vec![1],
162        param: b,
163        grad: None,
164    });
165    let q = g.mm(h, w);
166    let q = g.add(q, b);
167    g.reshape_(q, vec![batch as i64, 1])
168}
169
170fn build_train(spec: &RlSpec) -> CriticTrainGraph {
171    let f = DType::F32;
172    let b = spec.batch;
173    let sd = spec.state_dim;
174    let ad = spec.action_dim;
175    let mut g = Graph::new("critic_train");
176    let mut params = Vec::new();
177
178    let state = g.input("state", Shape::new(&[b, sd], f));
179    let action = g.input("action", Shape::new(&[b, ad], f));
180    let target = g.input("target", Shape::new(&[b, 1], f));
181
182    let feats = concat_features(&mut g, vec![state, action]);
183    let (h, in_d) = mlp_trunk(
184        &mut g,
185        feats,
186        spec.critic_in_dim(),
187        &spec.hidden,
188        "critic",
189        &mut params,
190    );
191    let q1 = q_head(&mut g, h, in_d, "critic_q1", b, &mut params);
192    let err = g.sub(q1, target);
193    let sq = g.mul(err, err);
194    let loss = g.mean(sq, vec![0, 1], false);
195    g.set_outputs(vec![loss]);
196    finalize_critic_train(g, params, loss)
197}
198
199fn build_q_grad(spec: &RlSpec) -> CriticQGradGraph {
200    let f = DType::F32;
201    let b = spec.batch;
202    let sd = spec.state_dim;
203    let ad = spec.action_dim;
204    let mut g = Graph::new("critic_q_grad");
205
206    let state = g.input("state", Shape::new(&[b, sd], f));
207    let action = g.input("action", Shape::new(&[b, ad], f));
208    let feats = concat_features(&mut g, vec![state, action]);
209    let mut params = Vec::new();
210    let (h, in_d) = mlp_trunk(
211        &mut g,
212        feats,
213        spec.critic_in_dim(),
214        &spec.hidden,
215        "critic",
216        &mut params,
217    );
218    let q1 = q_head(&mut g, h, in_d, "critic_q1", b, &mut params);
219    let loss = g.sum(q1, vec![0, 1], false);
220    g.set_outputs(vec![loss]);
221    let _ = params;
222
223    let (g, remap) = legalize_broadcast::run_with_remap(g);
224    let action = remap[&action];
225    let bwd = rlx_autodiff::grad_with_loss(&g, &[action]);
226    let d_output = bwd
227        .nodes()
228        .iter()
229        .find(|n| matches!(&n.op, Op::Input { name } if name == "d_output"))
230        .map(|n| n.id)
231        .expect("d_output");
232    let grad_action = bwd.outputs[1];
233
234    CriticQGradGraph {
235        graph: bwd,
236        d_output,
237        action,
238        grad_action,
239    }
240}
241
242fn finalize_critic_train(g: Graph, params: Vec<ParamSlot>, loss_fwd: NodeId) -> CriticTrainGraph {
243    let (g, remap) = legalize_broadcast::run_with_remap(g);
244    let params: Vec<ParamSlot> = params
245        .into_iter()
246        .map(|mut p| {
247            p.param = remap[&p.param];
248            p
249        })
250        .collect();
251    let loss_fwd = remap[&loss_fwd];
252    let wrt: Vec<NodeId> = params.iter().map(|p| p.param).collect();
253    let bwd = rlx_autodiff::grad_with_loss(&g, &wrt);
254    let d_output = bwd
255        .nodes()
256        .iter()
257        .find(|n| matches!(&n.op, Op::Input { name } if name == "d_output"))
258        .map(|n| n.id)
259        .expect("d_output");
260    let grad_ids: Vec<NodeId> = bwd.outputs[1..=params.len()].to_vec();
261    let params: Vec<ParamSlot> = params
262        .into_iter()
263        .zip(grad_ids)
264        .map(|(mut p, grad)| {
265            p.grad = Some(grad);
266            p
267        })
268        .collect();
269    CriticTrainGraph {
270        graph: bwd,
271        d_output,
272        loss: loss_fwd,
273        params,
274    }
275}
rlx_rl/graph/critic.rs

rlx_rl/graph/
critic.rs
