rename action_horizon to chunk_size

src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py

@@ -34,7 +34,7 @@ class PI05OpenPIConfig(PreTrainedConfig):
 
     # Input / output structure
     n_obs_steps: int = 1
-    action_horizon: int = 50  # Number of action steps to predict
+    chunk_size: int = 50  # Number of action steps to predict, in openpi called "action_horizon"
     n_action_steps: int = 50  # Number of action steps to execute
     action_dim: int = 32  # Action dimension (will be padded to 32)
     state_dim: int = 32  # State dimension (will be padded to 32)
@@ -87,9 +87,9 @@ class PI05OpenPIConfig(PreTrainedConfig):
         super().__post_init__()
 
         # Validate configuration
-        if self.n_action_steps > self.action_horizon:
+        if self.n_action_steps > self.chunk_size:
             raise ValueError(
-                f"n_action_steps ({self.n_action_steps}) cannot be greater than action_horizon ({self.action_horizon})"
+                f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
             )
 
         if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:

src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py

@@ -630,7 +630,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
         pad_masks.append(action_time_mask)
 
-        att_masks += [1] + ([0] * (self.config.action_horizon - 1))
+        att_masks += [1] + ([0] * (self.config.chunk_size - 1))
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
@@ -688,7 +688,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond
         )
 
-        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
         suffix_out = suffix_out.to(dtype=torch.float32)
 
         def action_out_proj_func(suffix_out):
@@ -711,7 +711,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         if noise is None:
             # Sample noise with padded dimension (32) as expected by action_in_proj
-            actions_shape = (bsize, self.config.action_horizon, 32)  # Use 32 for internal processing
+            actions_shape = (bsize, self.config.chunk_size, 32)  # Use 32 for internal processing
             noise = self.sample_noise(actions_shape, device)
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
@@ -789,7 +789,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         )
 
         suffix_out = outputs_embeds[1]
-        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
         suffix_out = suffix_out.to(dtype=torch.float32)
         return self.action_out_proj(suffix_out)
 

src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py

@@ -31,7 +31,7 @@ class PI0OpenPIConfig(PreTrainedConfig):
 
     # Input / output structure
     n_obs_steps: int = 1
-    action_horizon: int = 50  # Number of action steps to predict
+    chunk_size: int = 50  # Number of action steps to predict, in openpi called "action_horizon"
     n_action_steps: int = 50  # Number of action steps to execute
     action_dim: int = 32  # Action dimension (will be padded to 32)
     state_dim: int = 32  # State dimension (will be padded to 32)
@@ -84,9 +84,9 @@ class PI0OpenPIConfig(PreTrainedConfig):
         super().__post_init__()
 
         # Validate configuration
-        if self.n_action_steps > self.action_horizon:
+        if self.n_action_steps > self.chunk_size:
             raise ValueError(
-                f"n_action_steps ({self.n_action_steps}) cannot be greater than action_horizon ({self.action_horizon})"
+                f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
             )
 
         if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:

src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py

@@ -647,7 +647,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
         pad_masks.append(action_time_mask)
 
-        att_masks += [1] + ([0] * (self.config.action_horizon - 1))
+        att_masks += [1] + ([0] * (self.config.chunk_size - 1))
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
@@ -705,7 +705,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond
         )
 
-        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
         suffix_out = suffix_out.to(dtype=torch.float32)
 
         def action_out_proj_func(suffix_out):
@@ -728,7 +728,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         if noise is None:
             # Sample noise with padded dimension (32) as expected by action_in_proj
-            actions_shape = (bsize, self.config.action_horizon, 32)  # Use 32 for internal processing
+            actions_shape = (bsize, self.config.chunk_size, 32)  # Use 32 for internal processing
             noise = self.sample_noise(actions_shape, device)
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
@@ -806,7 +806,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         )
 
         suffix_out = outputs_embeds[1]
-        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out[:, -self.config.chunk_size :]
         suffix_out = suffix_out.to(dtype=torch.float32)
         return self.action_out_proj(suffix_out)