From 401bf431b8b0fe6e0a619412a607b5135dc4ded4 Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Wed, 1 Oct 2025 10:45:26 -0700
Subject: [PATCH] Modify Simulate::InjectNoise to set target of exponential
 decay to value from key_ctrl if a valid key is provided as input.

PiperOrigin-RevId: 813819040
Change-Id: I79508039b17fb3484b71ae0d19614c32adacfa10
---
 simulate/main.cc     | 4 ++--
 simulate/simulate.cc | 7 ++++++-
 simulate/simulate.h  | 4 ++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/simulate/main.cc b/simulate/main.cc
index 209be7e8..c2d72d95 100644
--- a/simulate/main.cc
+++ b/simulate/main.cc
@@ -381,7 +381,7 @@ void PhysicsLoop(mj::Simulate& sim) {
             sim.speed_changed = false;
 
             // inject noise
-            sim.InjectNoise();
+            sim.InjectNoise(sim.key);
 
             // run single step, let next iteration deal with timing
             mj_step(m, d);
@@ -412,7 +412,7 @@ void PhysicsLoop(mj::Simulate& sim) {
               }
 
               // inject noise
-              sim.InjectNoise();
+              sim.InjectNoise(sim.key);
 
               // call mj_step
               mj_step(m, d);
diff --git a/simulate/simulate.cc b/simulate/simulate.cc
index 555ab68c..3fb21976 100644
--- a/simulate/simulate.cc
+++ b/simulate/simulate.cc
@@ -2917,7 +2917,7 @@ void Simulate::AddToHistory() {
 }
 
 // inject Brownian noise
-void Simulate::InjectNoise() {
+void Simulate::InjectNoise(int key) {
   // no noise, return
   if (ctrl_noise_std <= 0) {
     return;
@@ -2936,6 +2936,11 @@ void Simulate::InjectNoise() {
       halfrange = 0.5 * (top - bottom);  // scales noise
     }
 
+    // overwrite midpoint with keyframe, if given
+    if (key >= 0) {
+      midpoint = m_->key_ctrl[key*m_->nu+i];
+    }
+
     // exponential convergence to midpoint at ctrl_noise_rate
     d_->ctrl[i] = rate * d_->ctrl[i] + (1-rate) * midpoint;
 
diff --git a/simulate/simulate.h b/simulate/simulate.h
index c1fd960e..567a1226 100644
--- a/simulate/simulate.h
+++ b/simulate/simulate.h
@@ -89,7 +89,7 @@ class Simulate {
   void AddToHistory();
 
   // inject control noise
-  void InjectNoise();
+  void InjectNoise(int key);
 
   // constants
   static constexpr int kMaxFilenameLength = 1000;
@@ -199,7 +199,7 @@ class Simulate {
   int busywait = 0;
 
   // keyframe index
-  int key = 0;
+  int key = -1;
 
   // index of history-scrubber slider
   int scrub_index = 0;