crates/training/rtx-transformers/src/layers/cloned_memory_updater.rs

@@ -294,12 +294,96 @@ impl ClonedMemoryUpdater {
         )
         .expect("backward");
 
+        // DAgger trains on the policy's own (clamped, sometimes extreme) visited
+        // states, which can spike the loss; a single non-finite step would
+        // poison every weight via Adam. Skip the update when the loss is not
+        // finite, and clip gradients so an extreme-but-finite state can't blow
+        // the optimizer up.
+        if !loss_val.is_finite() {
+            return loss_val;
+        }
         self.step += 1;
         let t = self.step;
         const LR: f32 = 0.01;
-        self.w_in.adam(&grad_of(&storage, in_id), LR, t);
-        self.w_mem.adam(&grad_of(&storage, mem_id), LR, t);
-        self.w_read.adam(&grad_of(&storage, read_id), LR, t);
+        const CLIP: f32 = 1.0;
+        let clip = |mut g: Vec<f32>| {
+            for v in &mut g {
+                *v = v.clamp(-CLIP, CLIP);
+            }
+            g
+        };
+        self.w_in.adam(&clip(grad_of(&storage, in_id)), LR, t);
+        self.w_mem.adam(&clip(grad_of(&storage, mem_id)), LR, t);
+        self.w_read.adam(&clip(grad_of(&storage, read_id)), LR, t);
+        loss_val
+    }
+
+    /// DAgger memory-only correction step: train just the recurrence
+    /// (`W_in`, `W_mem`) to map `(prev_mem, x) → target_mem` at learning rate
+    /// `lr`, leaving the readout `W_read` untouched. This is what on-policy
+    /// distillation needs — it pulls the free-running memory back toward the
+    /// oracle trajectory without disturbing the (BC-trained) readout, which
+    /// already works on oracle-range memory. Gradients are clipped and a
+    /// non-finite step is skipped, for stability on the extreme drifted states a
+    /// rollout visits. Returns the memory loss before the update.
+    ///
+    /// # Panics
+    /// Panics on any shape mismatch.
+    pub fn train_step_memory(
+        &mut self,
+        prev_mem: &[f32],
+        x: &[f32],
+        target_mem: &[f32],
+        lr: f32,
+    ) -> f32 {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.d_in;
+        assert_eq!(prev_mem.len(), dm);
+        assert_eq!(x.len(), din);
+        assert_eq!(target_mem.len(), dm);
+
+        let (w_in, in_id) = self.w_in.leaf(&self.dev);
+        let (w_mem, mem_id) = self.w_mem.leaf(&self.dev);
+
+        let mut cat = Vec::with_capacity(dm + din);
+        cat.extend_from_slice(prev_mem);
+        cat.extend_from_slice(x);
+        let cat_t = Ad::from_data(&cat, [1, dm + din], &self.dev);
+        let prev_t = Ad::from_data(prev_mem, [1, dm], &self.dev);
+
+        let h = Ad::gelu(Ad::matmul(cat_t, w_in));
+        let delta = Ad::matmul(h, w_mem);
+        let decay_t = Ad::from_data(&vec![DECAY; dm], [1, dm], &self.dev);
+        let new_mem = Ad::add(Ad::mul(prev_t, decay_t), delta);
+
+        let tgt_mem_t = Ad::from_data(target_mem, [1, dm], &self.dev);
+        let mem_err = Ad::sub(new_mem, tgt_mem_t);
+        let loss = Ad::sum(Ad::mul(mem_err.clone(), mem_err));
+        let loss_val = Ad::to_data(&loss)[0];
+
+        let storage = backward_impl(
+            &loss,
+            Some(GradTensor::from_d1(CpuBackend::ones(
+                [1],
+                &<CpuBackend as Backend>::Device::default(),
+            ))),
+        )
+        .expect("backward");
+
+        if !loss_val.is_finite() {
+            return loss_val;
+        }
+        self.step += 1;
+        let t = self.step;
+        const CLIP: f32 = 1.0;
+        let clip = |mut g: Vec<f32>| {
+            for v in &mut g {
+                *v = v.clamp(-CLIP, CLIP);
+            }
+            g
+        };
+        self.w_in.adam(&clip(grad_of(&storage, in_id)), lr, t);
+        self.w_mem.adam(&clip(grad_of(&storage, mem_id)), lr, t);
         loss_val
     }
 }