Struct NoGradTrack 

pub struct NoGradTrack { /* private fields */ }

RAII guard for temporarily disabling gradient tracking

NoGradTrack provides a scope-based mechanism for disabling gradient tracking, automatically restoring the previous state when the guard is dropped. It ensures proper gradient state management even in the presence of early returns or exceptions. A RAII guard that temporarily disables gradient tracking

Similar to PyTorch’s torch.no_grad(), this guard disables gradient computation within its scope and automatically restores the previous gradient tracking state when it goes out of scope.

Performance Benefits

• Prevents computation graph construction during inference
• Reduces memory usage by not storing intermediate values for backpropagation
• Improves computation speed by skipping gradient-related operations

Examples

use train_station::gradtrack::{NoGradTrack};
use train_station::Tensor;

let x = Tensor::ones(vec![3, 3]).with_requires_grad();
let y = Tensor::ones(vec![3, 3]).with_requires_grad();

// Normal computation with gradients
let z1 = x.add_tensor(&y);
assert!(z1.requires_grad());

// Computation without gradients
{
    let _guard = NoGradTrack::new();
    let z2 = x.add_tensor(&y);
    assert!(!z2.requires_grad()); // Gradients disabled
} // Guard drops here, gradients restored

// Gradients are automatically restored
let z3 = x.add_tensor(&y);
assert!(z3.requires_grad());

Nested Contexts

use train_station::{gradtrack::NoGradTrack, gradtrack::is_grad_enabled, Tensor};

assert!(is_grad_enabled());

{
    let _guard1 = NoGradTrack::new();
    assert!(!is_grad_enabled());

    {
        let _guard2 = NoGradTrack::new();
        assert!(!is_grad_enabled());
    } // guard2 drops

    assert!(!is_grad_enabled()); // Still disabled
} // guard1 drops

assert!(is_grad_enabled()); // Restored

Implementations

impl NoGradTrack

pub fn new() -> Self

Create a new NoGradTrack that disables gradient tracking

This function pushes the current gradient state onto the stack and disables gradient tracking. When the guard is dropped, the previous state is automatically restored.

Returns

A new NoGradTrack that will restore gradient state when dropped

Examples found in repository

examples/supervised_training/../neural_networks/feedforward_network.rs (line 61)

    pub fn forward_no_grad(input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        Self::forward(input)
    }
}

/// Configuration for feed-forward network
#[derive(Debug, Clone)]
pub struct FeedForwardConfig {
    pub input_size: usize,
    pub hidden_sizes: Vec<usize>,
    pub output_size: usize,
    #[allow(unused)]
    pub use_bias: bool,
}

impl Default for FeedForwardConfig {
    fn default() -> Self {
        Self {
            input_size: 4,
            hidden_sizes: vec![8, 4],
            output_size: 2,
            use_bias: true,
        }
    }
}

/// A configurable feed-forward neural network
pub struct FeedForwardNetwork {
    layers: Vec<LinearLayer>,
    config: FeedForwardConfig,
}

impl FeedForwardNetwork {
    /// Create a new feed-forward network with the given configuration
    pub fn new(config: FeedForwardConfig, seed: Option<u64>) -> Self {
        let mut layers = Vec::new();
        let mut current_size = config.input_size;
        let mut current_seed = seed;

        // Create hidden layers
        for &hidden_size in &config.hidden_sizes {
            layers.push(LinearLayer::new(current_size, hidden_size, current_seed));
            current_size = hidden_size;
            current_seed = current_seed.map(|s| s + 1);
        }

        // Create output layer
        layers.push(LinearLayer::new(
            current_size,
            config.output_size,
            current_seed,
        ));

        Self { layers, config }
    }

    /// Forward pass through the entire network
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let mut x = input.clone();

        // Pass through all layers except the last one with ReLU activation
        for layer in &self.layers[..self.layers.len() - 1] {
            x = layer.forward(&x);
            x = ReLU::forward(&x);
        }

        // Final layer without activation (raw logits)
        if let Some(final_layer) = self.layers.last() {
            x = final_layer.forward(&x);
        }

        x
    }

    /// Forward pass without gradients (for inference)
    #[allow(unused)]
    pub fn forward_no_grad(&self, input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        self.forward(input)
    }

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 81)

    pub fn forward_no_grad(&self, input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        self.forward(input)
    }

examples/RL_training/dqn.rs (line 310)

fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// Pseudo-Huber loss: sqrt(1 + diff^2) - 1 (smooth, robust)
fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/td3.rs (line 132)

    fn soft_update_from(&mut self, source: &Self, tau: f32) {
        let _ng = NoGradTrack::new();
        for (t, s) in self.layers.iter_mut().zip(source.layers.iter()) {
            // In-place Polyak update to preserve tensor IDs (no optimizer relink needed)
            let new_w = t
                .weight
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.weight.mul_scalar(tau));
            let new_b = t
                .bias
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.bias.mul_scalar(tau));
            {
                let src = new_w.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = new_b.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }
}

// -------------------------------
// Actor and Critic
// -------------------------------

struct Actor {
    net: Mlp,
}

impl Actor {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        // Smaller net for faster demo: sd -> 64 -> 64 -> ad, tanh output
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        Self { net }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state, Some(tanh_bounded))
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
    fn set_requires_grad_all(&mut self, enable: bool) {
        self.net.set_requires_grad_all(enable);
    }
}

struct Critic {
    net: Mlp,
}

impl Critic {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim + action_dim, 64, 64, 1], seed);
        Self { net }
    }
    fn forward(&self, state: &Tensor, action: &Tensor) -> Tensor {
        // Concatenate along feature dim (dim=1) for batched inputs
        // IMPORTANT: use views to preserve gradient graph; cloning would detach autograd
        let s_view = state.view(state.shape().dims().iter().map(|&d| d as i32).collect());
        let a_view = action.view(action.shape().dims().iter().map(|&d| d as i32).collect());
        let sa = Tensor::cat(&[s_view, a_view], 1);
        self.net.forward(&sa, None)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
    fn set_requires_grad_all(&mut self, enable: bool) {
        self.net.set_requires_grad_all(enable);
    }
}

// -------------------------------
// Simple continuous control environment: YardEnv
// State: normalized features [pos/3, clamp(vel/1, -1..1), bias(=0)] ; Action: scalar in [-1, 1]
// Dynamics: vel += 0.1*act - 0.01*pos; pos += vel
// Reward: -(pos^2) - 0.1*act^2 ; Episode ends if |pos| > 3 or step >= max_steps
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}

impl YardEnv {
    fn new(seed: u64) -> Self {
        let mut env = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        env.reset();
        env
    }

    fn reset(&mut self) -> Tensor {
        self.pos = self.rng.uniform(-0.5, 0.5);
        self.vel = self.rng.uniform(-0.1, 0.1);
        self.steps = 0;
        self.state_tensor()
    }

    fn state_tensor(&self) -> Tensor {
        // Normalize to keep critic inputs bounded:
        // - Position is bounded by termination at |pos|>3 → scale by 3 to [-1,1]
        // - Velocity scaled by 1.0 and clamped to [-1,1]
        let pos_n = self.pos / 3.0;
        let vel_n = self.vel.clamp(-1.0, 1.0);
        Tensor::from_slice(&[pos_n, vel_n, 0.0], vec![1, 3]).unwrap()
    }

    fn step(&mut self, action_value: f32) -> (Tensor, f32, bool) {
        let a = action_value.clamp(-1.0, 1.0);
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;

        let reward = -(self.pos * self.pos) - 0.1 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Replay Buffer
// -------------------------------

struct ReplayBuffer {
    capacity: usize,
    size: usize,
    pos: usize,
    state_dim: usize,
    action_dim: usize,
    states: Vec<f32>,
    actions: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    next_states: Vec<f32>,
}

impl ReplayBuffer {
    fn new(capacity: usize, state_dim: usize, action_dim: usize) -> Self {
        Self {
            capacity,
            size: 0,
            pos: 0,
            state_dim,
            action_dim,
            states: vec![0.0; capacity * state_dim],
            actions: vec![0.0; capacity * action_dim],
            rewards: vec![0.0; capacity],
            dones: vec![0.0; capacity],
            next_states: vec![0.0; capacity * state_dim],
        }
    }

    fn push(&mut self, s: &[f32], a: &[f32], r: f32, d: f32, s2: &[f32]) {
        let i = self.pos;
        let so = i * self.state_dim;
        let ao = i * self.action_dim;
        self.states[so..so + self.state_dim].copy_from_slice(s);
        self.actions[ao..ao + self.action_dim].copy_from_slice(a);
        self.rewards[i] = r;
        self.dones[i] = d;
        self.next_states[so..so + self.state_dim].copy_from_slice(s2);

        self.pos = (self.pos + 1) % self.capacity;
        self.size = self.size.saturating_add(1).min(self.capacity);
    }

    fn can_sample(&self, batch_size: usize) -> bool {
        self.size >= batch_size
    }

    fn sample(
        &self,
        batch_size: usize,
        rng: &mut SmallRng,
    ) -> (Tensor, Tensor, Tensor, Tensor, Tensor) {
        let mut s_vec = Vec::with_capacity(batch_size * self.state_dim);
        let mut a_vec = Vec::with_capacity(batch_size * self.action_dim);
        let mut r_vec = Vec::with_capacity(batch_size);
        let mut d_vec = Vec::with_capacity(batch_size);
        let mut s2_vec = Vec::with_capacity(batch_size * self.state_dim);

        for _ in 0..batch_size {
            let idx = rng.sample_index(self.size);
            let so = idx * self.state_dim;
            let ao = idx * self.action_dim;
            s_vec.extend_from_slice(&self.states[so..so + self.state_dim]);
            a_vec.extend_from_slice(&self.actions[ao..ao + self.action_dim]);
            r_vec.push(self.rewards[idx]);
            d_vec.push(self.dones[idx]);
            s2_vec.extend_from_slice(&self.next_states[so..so + self.state_dim]);
        }

        let s = Tensor::from_slice(&s_vec, vec![batch_size, self.state_dim]).unwrap();
        let a = Tensor::from_slice(&a_vec, vec![batch_size, self.action_dim]).unwrap();
        let r = Tensor::from_slice(&r_vec, vec![batch_size, 1]).unwrap();
        let d = Tensor::from_slice(&d_vec, vec![batch_size, 1]).unwrap();
        let s2 = Tensor::from_slice(&s2_vec, vec![batch_size, self.state_dim]).unwrap();
        (s, a, r, d, s2)
    }
}

// -------------------------------
// Helper: gradient clipping by global norm
// -------------------------------

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    // Compute global L2 norm of all grads
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                let scaled = g.mul_scalar(scale);
                p.set_grad(scaled);
            }
        }
    }
}

// Compute global L2 norm of gradients across a parameter list (read-only)
fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// Compute L2 norm of parameters (weights/biases) across a parameter list
fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main: TD3 training on YardEnv
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== TD3 Example (YardEnv) ===");

    // Environment / problem dims
    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hyperparameters (small for demo)
    let gamma = 0.99f32;
    let tau = 0.005f32; // Polyak
    let policy_noise = 0.2f32; // target smoothing noise stddev
    let exploration_noise = 0.1f32; // behavior policy noise stddev
    let policy_delay = 2usize;
    let batch_size = 64usize;
    let start_steps = 500usize; // random exploration steps
    let total_steps = 1500usize;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(11));
    let mut actor_targ = Actor::new(state_dim, action_dim, Some(12));
    actor_targ.net.copy_from(&actor.net);
    actor_targ.set_requires_grad_all(false);

    let mut critic1 = Critic::new(state_dim, action_dim, Some(21));
    let mut critic2 = Critic::new(state_dim, action_dim, Some(22));
    let mut critic1_targ = Critic::new(state_dim, action_dim, Some(23));
    let mut critic2_targ = Critic::new(state_dim, action_dim, Some(24));
    critic1_targ.net.copy_from(&critic1.net);
    critic2_targ.net.copy_from(&critic2.net);
    critic1_targ.set_requires_grad_all(false);
    critic2_targ.set_requires_grad_all(false);

    // Optimizers
    let mut actor_opt = Adam::with_learning_rate(1e-3);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }

    let mut critic_opt = Adam::with_learning_rate(1e-4);
    for p in critic1.parameters() {
        critic_opt.add_parameter(p);
    }
    for p in critic2.parameters() {
        critic_opt.add_parameter(p);
    }

    // Replay buffer and env
    let mut rb = ReplayBuffer::new(100_000, state_dim, action_dim);
    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(987654321);

    // Reset & metric trackers
    let mut state = env.reset(); // [1, state_dim]
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32; // smooth short-term
    let mut best_return = f32::NEG_INFINITY;
    let mut policy_updates: usize = 0;

    for t in 0..total_steps {
        // Select action
        let action_tensor = if t < start_steps {
            let a = rng.uniform(-1.0, 1.0);
            Tensor::from_slice(&[a], vec![1, action_dim]).unwrap()
        } else {
            // Behavior policy with exploration noise
            let _ng = NoGradTrack::new();
            let det = actor.forward(&state);
            let noise = Tensor::randn(vec![1, action_dim], None).mul_scalar(exploration_noise);
            tanh_bounded(&det.add_tensor(&noise))
        };
        let action_value = action_tensor.data()[0];

        // Environment step
        let (next_state, reward, done) = env.step(action_value);
        episode_return += reward;

        // Store transition
        let s_slice = state.data().to_vec();
        let a_slice = action_tensor.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            &a_slice,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        state = if done {
            let st = env.reset();
            // Metrics: update EMA and best
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={} | policy_updates={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size,
                policy_updates
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Training
        if rb.can_sample(batch_size) {
            // Sample batch
            let (s, a, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Compute target values y = r + (1-d)*gamma*min(Q1', Q2') using target networks (no grad)
            let target_q = {
                let _ng = NoGradTrack::new();
                // Target actions with smoothing noise (tanh bounds)
                let noise =
                    Tensor::randn(vec![batch_size, action_dim], None).mul_scalar(policy_noise);
                let a_targ = tanh_bounded(&actor_targ.forward(&s2).add_tensor(&noise));
                let q1_t = critic1_targ.forward(&s2, &a_targ);
                let q2_t = critic2_targ.forward(&s2, &a_targ);

                // Elementwise min via data() since this path is no-grad
                let q1d = q1_t.data();
                let q2d = q2_t.data();
                let mut min_vec = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let v1 = q1d[i];
                    let v2 = q2d[i];
                    min_vec.push(v1.min(v2));
                }
                let min_q = Tensor::from_slice(&min_vec, vec![batch_size, 1]).unwrap();
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&min_q))
            };

            // Critic update (both critics)
            // Zero grads in a short scope, then drop borrows before forward
            {
                let mut params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                critic_opt.zero_grad(&mut params);
            }

            // Forward current Q estimates
            let q1 = critic1.forward(&s, &a);
            let q2 = critic2.forward(&s, &a);
            let diff1 = q1.sub_tensor(&target_q);
            let diff2 = q2.sub_tensor(&target_q);
            let mut critic_loss = diff1
                .pow_scalar(2.0)
                .mean()
                .add_tensor(&diff2.pow_scalar(2.0).mean());

            // Backward
            critic_loss.backward(None);

            // Optional gradient clipping + step (only for params that received grads)
            {
                let params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    // Pre-step metrics
                    let grad_norm_before = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    critic_opt.step(&mut with_grads);
                    critic_opt.zero_grad(&mut with_grads);

                    // Post-step metrics (param norm)
                    let mut for_norm_params = {
                        let c_params = critic1.parameters();
                        let c2_params = critic2.parameters();
                        let mut tmp: Vec<&mut Tensor> = Vec::new();
                        tmp.extend(c_params);
                        tmp.extend(c2_params);
                        tmp
                    };
                    let param_norm = params_l2_norm(&mut for_norm_params);

                    // Print compact critic metrics occasionally
                    if t % 100 == 0 {
                        let q1_mean = q1.mean().value();
                        let q2_mean = q2.mean().value();
                        let tq_mean = target_q.mean().value();
                        println!(
                            "t={:5} | critic_loss={:.4} | q1_mean={:.3} q2_mean={:.3} tq_mean={:.3} | grad_norm={:.3} | crit_param_norm={:.3}",
                            t,
                            critic_loss.value(),
                            q1_mean,
                            q2_mean,
                            tq_mean,
                            grad_norm_before,
                            param_norm
                        );
                    }
                }
            }

            // Delayed policy update
            if t % policy_delay == 0 {
                // Actor update: maximize Q1(s, actor(s)) -> minimize -Q1
                // Zero actor grads before backward
                {
                    let mut a_params: Vec<&mut Tensor> = actor.parameters();
                    actor_opt.zero_grad(&mut a_params);
                }

                let a_pred = actor.forward(&s);
                let q_for_actor = critic1.forward(&s, &a_pred);
                let mut actor_loss = q_for_actor.mul_scalar(-1.0).mean();
                actor_loss.backward(None);

                {
                    let a_params: Vec<&mut Tensor> = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in a_params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let grad_norm_before = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);

                        // Post-step param norm
                        let mut for_norm_params = actor.parameters();
                        let param_norm = params_l2_norm(&mut for_norm_params);

                        policy_updates += 1;
                        if t % 200 == 0 {
                            println!(
                                "t={:5} | actor_loss={:.4} | act_grad_norm={:.3} | act_param_norm={:.3} | lr_a={:.4e} lr_c={:.4e} | policy_updates={}",
                                t,
                                actor_loss.value(),
                                grad_norm_before,
                                param_norm,
                                actor_opt.learning_rate(),
                                critic_opt.learning_rate(),
                                policy_updates
                            );
                        }
                    }
                }

                // Target updates (Polyak averaging, no grad)
                actor_targ.net.soft_update_from(&actor.net, tau);
                critic1_targ.net.soft_update_from(&critic1.net, tau);
                critic2_targ.net.soft_update_from(&critic2.net, tau);
            }

            // Clear entire graphs to avoid stale accumulation across iterations
            clear_all_graphs_known();
        }
    }

    println!("=== TD3 training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 378)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

Trait Implementations

impl Default for NoGradTrack

fn default() -> Self

Returns the “default value” for a type.

impl Drop for NoGradTrack

fn drop(&mut self)

Automatically restore the previous gradient tracking state

This ensures that gradient tracking is properly restored even if the guard goes out of scope due to early returns or panics.