fix(rollout): fixing the episode length to use hwc (#3469)

also reducing default length to 5 minutes
This commit is contained in:
Maxime Ellerbach
2026-04-27 15:42:46 +02:00
committed by GitHub
parent c505a6b1f4
commit bef67960d0
+9 -9
View File
@@ -219,7 +219,7 @@ def estimate_max_episode_seconds(
The estimate ignores codec-specific settings (CRF, preset) on purpose: The estimate ignores codec-specific settings (CRF, preset) on purpose:
we only need a rough lower bound on bitrate, not a precise prediction. we only need a rough lower bound on bitrate, not a precise prediction.
Falls back to 600 s (10 min) when no video features are present. Falls back to 300 s (5 min) when no video features are present.
""" """
# 0.1 bits-per-pixel is a *low* estimate for CRF-30 streaming video of # 0.1 bits-per-pixel is a *low* estimate for CRF-30 streaming video of
# robot footage (real-world is typically 0.1 0.3 bpp). Under- # robot footage (real-world is typically 0.1 0.3 bpp). Under-
@@ -233,16 +233,16 @@ def estimate_max_episode_seconds(
if feat.get("dtype") == "video": if feat.get("dtype") == "video":
shape = feat.get("shape", ()) shape = feat.get("shape", ())
# Assuming shape could be (C, H, W) or (T, C, H, W) # (H, W, C)
# We want to extract the spatial dimensions. # We want to extract the spatial dimensions.
if len(shape) >= 3: if len(shape) == 3:
h, w = shape[-2], shape[-1] pixels = shape[0] * shape[1] * shape[2]
pixels = h * w camera_pixels.append(pixels)
if pixels > 0: else:
camera_pixels.append(pixels) raise ValueError(f"Unexpected video feature shape: {shape}")
if not camera_pixels: if not camera_pixels:
return 600.0 return 300.0
# Use the smallest camera: it produces the lowest bitrate and therefore # Use the smallest camera: it produces the lowest bitrate and therefore
# takes the longest to reach the target — the conservative choice. # takes the longest to reach the target — the conservative choice.
@@ -252,7 +252,7 @@ def estimate_max_episode_seconds(
# Guard against division by zero just in case # Guard against division by zero just in case
if bytes_per_second <= 0: if bytes_per_second <= 0:
return 600.0 return 300.0
return (target_size_mb * 1024 * 1024) / bytes_per_second return (target_size_mb * 1024 * 1024) / bytes_per_second