[Serving] Update the details of the inference service (#70)

* Update the details of the inference service startup

* add functions to compute and update robot-specific action/state normalization stats

* Update for pre-commit

* Union compute_action_statistics code into update_action_statistics

* Update for pre-commit

* Add default value for dof_config and agent_pos_config

---------

Co-authored-by: farmer <farmer@x2robot.com>

wall_x/data/utils.py

@@ -5,14 +5,16 @@ This module provides utilities for preprocessing text, images, and actions
 for multimodal transformer models in robotic learning tasks.
 """
 
-import re
-import torch
-import random
-from collections import OrderedDict
-from typing import List, Dict, Any, Optional, Union, Tuple
-from transformers import BatchFeature
-from dataclasses import dataclass
 import json
+import random
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+
+import torch
+from transformers import BatchFeature
 
 KEY_MAPPINGS = {
     "lerobot/aloha_mobile_cabinet": {
@@ -32,6 +34,15 @@ KEY_MAPPINGS = {
         "state": "state",
         "action": "actions",
     },
+    "x2": {
+        "camera": {
+            "observation.images.faceImg": "face_view",
+            "observation.images.leftImg": "left_wrist_view",
+            "observation.images.rightImg": "right_wrist_view",
+        },
+        "state": "state",
+        "action": "actions",
+    },
 }
 
 CAMERA_NAME_MAPPING = {
@@ -662,3 +673,85 @@ def load_norm_stats(norm_stats_path, dataset_name):
     )
 
     return {"action": action_norm_stats, "state": state_norm_stats}
+
+
+def update_action_statistics(
+    action_statistic_dof: Dict[str, Any],
+    norm_stats_path: str,
+    repo_id: str,
+    dof_config: Dict[str, int] = None,
+    agent_pos_config: Dict[str, int] = None,
+    robot_name: str = None,
+    customized_dof_config: Dict[str, int] = None,
+    customized_agent_pos_config: Dict[str, int] = None,
+) -> None:
+    """
+    Update the action statistics dictionary with new robot configuration.
+
+    Args:
+        action_statistic_dof (Dict[str, Any]): The dictionary to be updated with statistics
+        norm_stats_path (str): Path to the normalization statistics file
+        repo_id (str): Repository ID for the LeRobot configuration
+        dof_config (Dict[str, int]): Configuration mapping DOF names to their dimensions
+        agent_pos_config (Dict[str, int]): Configuration mapping agent position names to their dimensions
+        robot_name (str, optional): Name of the robot. If None, uses repo_id as the key
+        customized_dof_config (Dict[str, int], optional): Customized DOF configuration for specific robot
+        customized_agent_pos_config (Dict[str, int], optional): Customized agent position configuration for specific robot
+    """
+    # Load normalization statistics
+    norm_stats = load_norm_stats(norm_stats_path, repo_id)
+
+    # Extract min and delta values for action and state
+    action_min = norm_stats["action"].min.numpy().tolist()
+    action_delta = norm_stats["action"].delta.numpy().tolist()
+    state_min = norm_stats["state"].min.numpy().tolist()
+    state_delta = norm_stats["state"].delta.numpy().tolist()
+
+    # Use customized configurations if provided, otherwise use default ones
+    current_dof_config = (
+        customized_dof_config if customized_dof_config is not None else dof_config
+    )
+    current_agent_pos_config = (
+        customized_agent_pos_config
+        if customized_agent_pos_config is not None
+        else agent_pos_config
+    )
+
+    # Prepare keys and values for DOF and agent position configurations
+    dof_key = []
+    agent_pos_key = []
+    dof_value = []
+    agent_pos_value = []
+    stats_dict = {}
+
+    # Extract DOF configuration
+    for k, v in current_dof_config.items():
+        dof_key.append(k)
+        dof_value.append(v)
+
+    # Extract agent position configuration
+    for k, v in current_agent_pos_config.items():
+        agent_pos_key.append(k)
+        agent_pos_value.append(v)
+
+    # Calculate DOF indices and extract corresponding min/delta values
+    dof_idx = np.array([0] + dof_value).cumsum()
+    for i in range(len(dof_idx) - 1):
+        stats_dict[dof_key[i]] = {
+            "min": action_min[dof_idx[i] : dof_idx[i + 1]],
+            "delta": action_delta[dof_idx[i] : dof_idx[i + 1]],
+        }
+
+    # Calculate agent position indices and extract corresponding min/delta values
+    agent_pos_idx = np.array([0] + agent_pos_value).cumsum()
+    for i in range(len(agent_pos_idx) - 1):
+        stats_dict[agent_pos_key[i]] = {
+            "min": state_min[agent_pos_idx[i] : agent_pos_idx[i + 1]],
+            "delta": state_delta[agent_pos_idx[i] : agent_pos_idx[i + 1]],
+        }
+
+    # Use provided robot name or repo_id as the key
+    robot_key = robot_name if robot_name is not None else repo_id
+
+    # Update the action_statistic_dof dictionary
+    action_statistic_dof.update({robot_key: stats_dict})