chore: replace hard-coded action values with constants throughout all the source code (#2055)

* chore: replace hard-coded 'action' values with constants throughout all the source code

* chore(tests): replace hard-coded action values with constants throughout all the test code

src/lerobot/policies/act/modeling_act.py

@@ -394,7 +394,7 @@ class ACT(nn.Module):
             latent dimension.
         """
         if self.config.use_vae and self.training:
-            assert "action" in batch, (
+            assert ACTION in batch, (
                 "actions must be provided when using the variational objective in training mode."
             )
 
@@ -404,7 +404,7 @@ class ACT(nn.Module):
             batch_size = batch[OBS_ENV_STATE].shape[0]
 
         # Prepare the latent for input to the transformer encoder.
-        if self.config.use_vae and "action" in batch and self.training:
+        if self.config.use_vae and ACTION in batch and self.training:
             # Prepare the input to the VAE encoder: [cls, *joint_space_configuration, *action_sequence].
             cls_embed = einops.repeat(
                 self.vae_encoder_cls_embed.weight, "1 d -> b 1 d", b=batch_size
@@ -412,7 +412,7 @@ class ACT(nn.Module):
             if self.config.robot_state_feature:
                 robot_state_embed = self.vae_encoder_robot_state_input_proj(batch[OBS_STATE])
                 robot_state_embed = robot_state_embed.unsqueeze(1)  # (B, 1, D)
-            action_embed = self.vae_encoder_action_input_proj(batch["action"])  # (B, S, D)
+            action_embed = self.vae_encoder_action_input_proj(batch[ACTION])  # (B, S, D)
 
             if self.config.robot_state_feature:
                 vae_encoder_input = [cls_embed, robot_state_embed, action_embed]  # (B, S+2, D)

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -82,7 +82,7 @@ class DiffusionPolicy(PreTrainedPolicy):
         """Clear observation and action queues. Should be called on `env.reset()`"""
         self._queues = {
             OBS_STATE: deque(maxlen=self.config.n_obs_steps),
-            "action": deque(maxlen=self.config.n_action_steps),
+            ACTION: deque(maxlen=self.config.n_action_steps),
         }
         if self.config.image_features:
             self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps)
@@ -306,10 +306,10 @@ class DiffusionModel(nn.Module):
         }
         """
         # Input validation.
-        assert set(batch).issuperset({OBS_STATE, "action", "action_is_pad"})
+        assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"})
         assert OBS_IMAGES in batch or OBS_ENV_STATE in batch
         n_obs_steps = batch[OBS_STATE].shape[1]
-        horizon = batch["action"].shape[1]
+        horizon = batch[ACTION].shape[1]
         assert horizon == self.config.horizon
         assert n_obs_steps == self.config.n_obs_steps
 
@@ -317,7 +317,7 @@ class DiffusionModel(nn.Module):
         global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)
 
         # Forward diffusion.
-        trajectory = batch["action"]
+        trajectory = batch[ACTION]
         # Sample noise to add to the trajectory.
         eps = torch.randn(trajectory.shape, device=trajectory.device)
         # Sample a random noising timestep for each item in the batch.
@@ -338,7 +338,7 @@ class DiffusionModel(nn.Module):
         if self.config.prediction_type == "epsilon":
             target = eps
         elif self.config.prediction_type == "sample":
-            target = batch["action"]
+            target = batch[ACTION]
         else:
             raise ValueError(f"Unsupported prediction type {self.config.prediction_type}")
 

src/lerobot/policies/pi0/conversion_scripts/compare_with_jax.py

@@ -21,7 +21,7 @@ import torch
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.policies.factory import make_policy
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
+from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 
 def display(tensor: torch.Tensor):
@@ -73,7 +73,7 @@ def main():
     for cam_key, uint_chw_array in example["images"].items():
         batch[f"{OBS_IMAGES}.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
     batch[OBS_STATE] = torch.from_numpy(example["state"])
-    batch["action"] = torch.from_numpy(outputs["actions"])
+    batch[ACTION] = torch.from_numpy(outputs["actions"])
     batch["task"] = example["prompt"]
 
     if model_name == "pi0_aloha_towel":
@@ -117,7 +117,7 @@ def main():
         actions.append(action)
 
     actions = torch.stack(actions, dim=1)
-    pi_actions = batch["action"]
+    pi_actions = batch[ACTION]
     print("actions")
     display(actions)
     print()

src/lerobot/policies/sac/configuration_sac.py

@@ -225,7 +225,7 @@ class SACConfig(PreTrainedConfig):
                 "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
             )
 
-        if "action" not in self.output_features:
+        if ACTION not in self.output_features:
             raise ValueError("You must provide 'action' in the output features")
 
     @property

src/lerobot/policies/sac/modeling_sac.py

@@ -31,7 +31,7 @@ from torch.distributions import MultivariateNormal, TanhTransform, Transform, Tr
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_STATE
+from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
 
 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension
 
@@ -51,7 +51,7 @@ class SACPolicy(
         self.config = config
 
         # Determine action dimension and initialize all components
-        continuous_action_dim = config.output_features["action"].shape[0]
+        continuous_action_dim = config.output_features[ACTION].shape[0]
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -158,7 +158,7 @@ class SACPolicy(
             The computed loss tensor
         """
         # Extract common components from batch
-        actions: Tensor = batch["action"]
+        actions: Tensor = batch[ACTION]
         observations: dict[str, Tensor] = batch["state"]
         observation_features: Tensor = batch.get("observation_feature")
 

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -92,7 +92,7 @@ class TDMPCPolicy(PreTrainedPolicy):
         """
         self._queues = {
             OBS_STATE: deque(maxlen=1),
-            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
+            ACTION: deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
         }
         if self.config.image_features:
             self._queues[OBS_IMAGE] = deque(maxlen=1)