chore: use alias & constants (#2785)

* chore: use alias and constants

* fix(rl): solve circular dependecy

* chore: nit right constant

* chore: pre-commit

* chore(script): conflict tokenizer train

---------

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>

src/lerobot/utils/constants.py

@@ -28,11 +28,13 @@ OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
 
 ACTION = "action"
+ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"
 ACTION_TOKEN_MASK = ACTION + ".token_mask"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"
+INFO = "info"
 
 ROBOTS = "robots"
 TELEOPERATORS = "teleoperators"

src/lerobot/utils/visualization_utils.py

@@ -14,12 +14,13 @@
 
 import numbers
 import os
-from typing import Any
 
 import numpy as np
 import rerun as rr
 
-from .constants import OBS_PREFIX, OBS_STR
+from lerobot.processor import RobotAction, RobotObservation
+
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
 
 
 def init_rerun(
@@ -50,8 +51,8 @@ def _is_scalar(x):
 
 
 def log_rerun_data(
-    observation: dict[str, Any] | None = None,
-    action: dict[str, Any] | None = None,
+    observation: RobotObservation | None = None,
+    action: RobotAction | None = None,
     compress_images: bool = False,
 ) -> None:
     """
@@ -96,7 +97,7 @@ def log_rerun_data(
         for k, v in action.items():
             if v is None:
                 continue
-            key = k if str(k).startswith("action.") else f"action.{k}"
+            key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}"
 
             if _is_scalar(v):
                 rr.log(key, rr.Scalars(float(v)))