From 16788f847ed2cd710d798d795ae5ad7339a55070 Mon Sep 17 00:00:00 2001
From: Karl Pertsch <karl@physicalintelligence.company>
Date: Sat, 8 Feb 2025 23:33:52 +0000
Subject: [PATCH] cap DROID execution frequency

---
 examples/droid/main.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/examples/droid/main.py b/examples/droid/main.py
index 87dc928..93b7654 100644
--- a/examples/droid/main.py
+++ b/examples/droid/main.py
@@ -6,7 +6,7 @@ import datetime
 import faulthandler
 import os
 import signal
-
+import time
 from moviepy.editor import ImageSequenceClip
 import numpy as np
 from openpi_client import image_tools
@@ -19,6 +19,9 @@ import tyro
 
 faulthandler.enable()
 
+# DROID data collection frequency -- we slow down execution to match this frequency
+DROID_CONTROL_FREQUENCY = 15
+
 
 @dataclasses.dataclass
 class Args:
@@ -95,6 +98,7 @@ def main(args: Args):
         bar = tqdm.tqdm(range(args.max_timesteps))
         print("Running rollout... press Ctrl+C to stop early.")
         for t_step in bar:
+            start_time = time.time()
             try:
                 # Get the current observation
                 curr_obs = _extract_observation(
@@ -145,6 +149,11 @@ def main(args: Args):
                 action = np.clip(action, -1, 1)
 
                 env.step(action)
+
+                # Sleep to match DROID data collection frequency
+                elapsed_time = time.time() - start_time
+                if elapsed_time < 1 / DROID_CONTROL_FREQUENCY:
+                    time.sleep(1 / DROID_CONTROL_FREQUENCY - elapsed_time)
             except KeyboardInterrupt:
                 break