Module tutorial_07_wasm 

Tutorial 7 — Train in Rust, run in the browser.

See the rendered Markdown at docs/tutorials/07-wasm-deploy.md.

Tutorial 7 — Train in Rust, run in the browser

You will: take a policy trained by any earlier tutorial and export it to the JSON inference format that Thrust’s WebAssembly demo consumes — including the one weight-layout gotcha that silently breaks a naive export — then wire the JSON into the running browser demo without touching a line of TypeScript. Prerequisites: any tutorial that produces a trained MlpBurnPolicy — Tutorial 2 (CartPole + PPO) is the natural pairing, since the demo ships a CartPole page. You need to recognize MlpBurnPolicy and the MlpBurnConfig surface; nothing about the browser stack is assumed. Time: ~15 minutes.

This is the last tutorial in the series, and it closes the loop the first six opened: you have trained policies — now you ship one. The destination is the live WebAssembly demo under web/, where a CartPole, Pong, or Snake agent runs its forward pass entirely in the browser, at 60 FPS, with no server round-trip.

The thing to understand before any code is why deployment needs a second model format at all.

Why a separate inference format

Training and inference in Thrust deliberately run on two different stacks:

• Training uses Burn, a full tensor framework with autodiff, optimizers, and backend-swappable kernels (NdArray on CPU, WGPU on GPU). That machinery is exactly what you want while learning — and exactly what you do not want to compile to WebAssembly, because it drags a large, GPU-flavored dependency tree into a bundle that has to download over the wire.
• Inference in the browser uses a tiny, pure-Rust struct — InferenceModel — that holds the weights as plain Vec<Vec<f32>> matrices and does the forward pass with nested for loops. No Burn, no tensors, no autodiff: it compiles to a small WASM module and runs a CartPole forward pass in well under a millisecond.

The bridge between the two stacks is a JSON file. You train with Burn, export the learned weights into an InferenceModel, serialize it to JSON, and the WASM demo deserializes that JSON at page load. The JSON is the contract: as long as it has the right field names and shapes, the browser never needs to know Burn existed.

A note on InferenceModel vs ExportedModel. The crate also contains a crate::inference::ExportedModel type — a simpler, single-layer format that is not what the WASM demo loads. The bindings in src/wasm.rs (WasmCartPole, WasmPong, WasmSnake) all deserialize into thrust_rl::policy::inference::InferenceModel. That is the type this tutorial targets, and the type whose field names you will see in web/public/cartpole_model_best.json.

The inference format, field by field

InferenceModel mirrors a 2-layer shared-trunk actor-critic — the architecture MlpBurnConfig { num_layers: 2, .. } builds:

InferenceModel {
    obs_dim:    usize,   // network input width  (CartPole: 4)
    action_dim: usize,   // policy-head output   (CartPole: 2)
    hidden_dim: usize,   // trunk width          (e.g. 64)
    activation: InferenceActivation,  // "Tanh" or "ReLU", applied on the trunk

    // Trunk (shared between policy and value heads)
    shared_fc1_weight: Vec<Vec<f32>>,  // [hidden_dim × obs_dim]
    shared_fc1_bias:   Vec<f32>,       // [hidden_dim]
    shared_fc2_weight: Vec<Vec<f32>>,  // [hidden_dim × hidden_dim]
    shared_fc2_bias:   Vec<f32>,       // [hidden_dim]

    // Heads
    policy_weight: Vec<Vec<f32>>,      // [action_dim × hidden_dim]  -> action logits
    policy_bias:   Vec<f32>,           // [action_dim]
    value_weight:  Vec<Vec<f32>>,      // [1 × hidden_dim]           -> scalar value
    value_bias:    Vec<f32>,           // [1]

    metadata: Option<TrainingMetadata>, // provenance; optional
}

Every weight matrix is stored out-major: weight[o][i] is the connection from input neuron i to output neuron o, so each row is one output neuron’s full set of incoming weights. That row-major convention is what InferenceModel::forward iterates over, and it is the crux of the one gotcha in this tutorial.

The weight-layout gotcha

Burn’s Linear layer stores its weight tensor with shape [in_features, out_features] — the transpose of what InferenceModel wants. If you copy the flat Burn buffer straight into an InferenceModel row, every weight lands in the wrong cell and the exported policy plays like noise, even though the JSON is “valid” and loads without error. This is the single most common way a WASM export goes silently wrong.

The fix is a transpose during the copy. Burn’s flat buffer indexes as flat[i * out_features + o] (input-major); InferenceModel wants rows[o][i] (output-major). So:

for o in 0..out_features {
    for i in 0..in_features {
        rows[o][i] = flat[i * out_features + o];
    }
}

The reference implementation lives in extract_linear() in examples/games/pong/eval_pong.rs, which is what actually produced the shipped Pong weights. The doc-test below is the same transpose, distilled to the essentials.

The code

The snippet below is the complete export path, and it is a doc-test — it compiles and runs against the current API on every CI run, so the transpose can never silently rot out of sync with Burn’s tensor layout.

One deliberate shortcut: it does not train. In production you train first (that is what Tutorials 2–5 are for) and then export; a real training run takes minutes to hours, which is far too slow for a doc-test. Since the export machinery — layer extraction, transpose, JSON round-trip — is completely independent of what the weights are, we exercise it on a freshly initialized (random) policy. The structure is identical; only the numbers would differ after training. Copy it into a fn main() (and cargo add thrust-rl, plus burn and tempfile) and it runs unchanged.

use burn::{
    backend::{Autodiff, NdArray},
    nn::Linear,
};
use thrust_rl::policy::inference::{InferenceActivation, InferenceModel};
use thrust_rl::prelude::*;

type Backend = Autodiff<NdArray<f32>>;

// --- Architecture (must match the trained policy you are exporting) --------
const OBS_DIM: usize = 4; // CartPole observation width
const ACTION_DIM: usize = 2; // CartPole: push left / right
const HIDDEN_DIM: usize = 64; // trunk width

let device = Default::default();

// --- The policy ------------------------------------------------------------
// In a real workflow this is the policy you trained in Tutorial 2 and loaded
// back from a Burn `.bin` record. Here we build a fresh, randomly initialized
// one: the export path is identical either way, and a doc-test cannot afford a
// training run. `Tanh` matches the activation the shipped CartPole model uses.
let config = MlpBurnConfig {
    num_layers: 2,
    hidden_dim: HIDDEN_DIM,
    use_orthogonal_init: true,
    activation: BurnActivation::Tanh,
    seed: Some(0),
};
let policy = MlpBurnPolicy::<Backend>::with_config(OBS_DIM, ACTION_DIM, config, &device);

// --- Extract one Burn Linear layer as an out-major (weight, bias) pair -----
// Burn stores the weight tensor as [in_features, out_features] in a flat,
// input-major buffer. InferenceModel wants out-major rows, weight[out][in], so
// we transpose during the copy. Get this wrong and the JSON still "loads" — it
// just plays like noise.
fn extract_linear(layer: &Linear<Backend>) -> (Vec<Vec<f32>>, Vec<f32>) {
    let weight = layer.weight.val();
    let [in_features, out_features] = weight.dims();
    let flat: Vec<f32> = weight.into_data().to_vec().expect("weight to_vec");

    let mut rows = vec![vec![0.0_f32; in_features]; out_features];
    for o in 0..out_features {
        for i in 0..in_features {
            // Burn flat index: flat[i * out_features + o]  (in = i, out = o).
            rows[o][i] = flat[i * out_features + o];
        }
    }

    // A Linear may or may not carry a bias; default to zeros when it doesn't.
    let bias = layer
        .bias
        .as_ref()
        .map(|b| b.val().into_data().to_vec::<f32>().expect("bias to_vec"))
        .unwrap_or_else(|| vec![0.0_f32; out_features]);

    (rows, bias)
}

// --- Build the InferenceModel from the four extracted layers ---------------
// The layer accessors (fc1/fc2/policy_head/value_head) expose exactly the four
// Linear layers of the shared-trunk actor-critic.
let (shared_fc1_weight, shared_fc1_bias) = extract_linear(policy.fc1());
let (shared_fc2_weight, shared_fc2_bias) = extract_linear(policy.fc2());
let (policy_weight, policy_bias) = extract_linear(policy.policy_head());
let (value_weight, value_bias) = extract_linear(policy.value_head());

let model = InferenceModel {
    obs_dim: OBS_DIM,
    action_dim: ACTION_DIM,
    hidden_dim: HIDDEN_DIM,
    activation: InferenceActivation::Tanh,
    metadata: None, // set TrainingMetadata here to record provenance
    shared_fc1_weight,
    shared_fc1_bias,
    shared_fc2_weight,
    shared_fc2_bias,
    policy_weight,
    policy_bias,
    value_weight,
    value_bias,
};

// Sanity-check the shapes before serializing: each weight is out-major, so its
// row count is the output width and each row's length is the input width.
assert_eq!(model.shared_fc1_weight.len(), HIDDEN_DIM); // [hidden × obs]
assert_eq!(model.shared_fc1_weight[0].len(), OBS_DIM);
assert_eq!(model.policy_weight.len(), ACTION_DIM); // [action × hidden]
assert_eq!(model.policy_weight[0].len(), HIDDEN_DIM);
assert_eq!(model.value_weight.len(), 1); // [1 × hidden]

// --- Round-trip through JSON, exactly as the browser will ------------------
// `save_json` writes the same wire format the WASM demo fetches; `load_json`
// parses it back. We use a temp file so the doc-test leaves nothing behind; in
// practice you write straight to web/public/<name>.json (see below).
let file = tempfile::NamedTempFile::new().expect("temp file");
model.save_json(file.path()).expect("save_json");
let loaded = InferenceModel::load_json(file.path()).expect("load_json");

// The four fields the demo keys on survive the round-trip intact.
assert_eq!(loaded.obs_dim, OBS_DIM);
assert_eq!(loaded.action_dim, ACTION_DIM);
assert_eq!(loaded.hidden_dim, HIDDEN_DIM);
assert!(matches!(loaded.activation, InferenceActivation::Tanh));

// And the loaded model actually runs a forward pass: obs -> (logits, value).
let (logits, value) = loaded.forward(&[0.0, 0.0, 0.0, 0.0]);
assert_eq!(logits.len(), ACTION_DIM);
assert!(value.is_finite());

That is the entire native side of deployment: extract, transpose, build, serialize. The JSON save_json produces has the top-level keys obs_dim, action_dim, hidden_dim, activation, the six weight/bias fields, and an optional metadata block — byte-for-byte the structure of the shipped web/public/cartpole_model_best.json.

Wiring it into the browser

The WASM demo is already fully wired — there is no TypeScript to write. The CartPole page fetches its weights by a fixed filename at startup and hands the raw JSON to the WASM binding:

// web/src/components/CartPole/useCartPole.ts (already in the repo)
const response = await fetch(`${import.meta.env.BASE_URL}cartpole_model_best.json`);
const modelJson = await response.text();
envRef.current.load_policy_json(modelJson); // WasmCartPole::load_policy_json

load_policy_json on the Rust side (src/wasm.rs) is just InferenceModel::load_json over a string — the same deserialize the doc-test round-trips. So “deploying your policy” means putting your JSON where that fetch looks for it. End to end:

# 1. Export your trained policy to the filename the demo already expects.
#    (In your own binary, call model.save_json("web/public/cartpole_model_best.json").)

# 2. Build the WASM bundle from the crate root. --no-default-features drops the
#    Burn training stack; the `wasm` feature pulls in only the inference path.
wasm-pack build --target web --features wasm --no-default-features

# 3. Start the demo's dev server.
cd web && pnpm install && pnpm dev

Then open the printed local URL. On load the CartPole page fetches cartpole_model_best.json, calls load_policy_json, and your exported policy starts balancing the pole — running its forward pass in the browser, no server involved. To swap in a different policy, overwrite that JSON and reload; the filename is the only contract. Pong and Snake follow the same pattern with their own fixed filenames (pong_model.json, snake_model.json) and bindings (WasmPong, WasmSnake).

See web/README.md for the full demo setup and docs/WASM_ROADMAP.md for the broader train-export-deploy pipeline.

Try it yourself

• Export a trained policy. Run Tutorial 2 to convergence, save the Burn record, then load it into a blank MlpBurnPolicy::<Backend>::with_config(..) with load_file (see eval_pong.rs) before the extraction step. The export code above does not change — only the weights it reads.
• Attach provenance. Fill in metadata: Some(TrainingMetadata { .. }) with the environment name, algorithm, and step count. It rides along in the JSON and shows up when you inspect a shipped model.
• Match the activation. If you trained with BurnActivation::ReLU, export with InferenceActivation::ReLU — the field controls which nonlinearity InferenceModel::forward applies. A mismatch here is another silent correctness bug, exactly like the transpose.
• Verify parity. Before and after export, feed the same observation through the Burn policy and the InferenceModel and compare the logits. Agreement (to float tolerance) proves your transpose and activation are right — the surest check that the browser will behave like training did.

Next

That closes the series: you can install Thrust, train a policy with PPO, DQN, SAC, or recurrent PPO, write your own environment, and now take a trained policy all the way to a browser tab. From here, the Example Gallery has full runnable trainers for every algorithm, and docs/WASM_ROADMAP.md covers the deployment pipeline in depth. See the tutorial index for the whole path at a glance.