diff --git a/scripts/v2_rl_walk_mujoco.py b/scripts/v2_rl_walk_mujoco.py
index 6e6e569..931420e 100644
--- a/scripts/v2_rl_walk_mujoco.py
+++ b/scripts/v2_rl_walk_mujoco.py
@@ -261,7 +261,7 @@ class RLWalk:
                 if obs is None:
                     continue
 
-                if np.linalg.norm(self.last_commands[3:]) > 0.01:
+                if np.linalg.norm(self.last_commands[:3]) > 0.01:
                     self.imitation_i += 1 * (
                         self.phase_frequency_factor + self.phase_frequency_factor_offset
                     )