update recent changes (#58)

Co-authored-by: yuanhaonan <yuanhaonan@dexforce.top>
Co-authored-by: yuecideng <dengyueci@qq.com>

Author: 3 people

embodichain/lab/devices/__init__.py

@@ -15,3 +15,9 @@
 # ----------------------------------------------------------------------------
 
 from .device import Device
+from .device_controller import DeviceController
+
+__all__ = [
+    "Device",
+    "DeviceController",
+]

embodichain/lab/devices/device_controller.py

@@ -0,0 +1,295 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2025 DexForce Technology Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ----------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import torch
+from typing import TYPE_CHECKING, Dict, Any, Optional, List, Union
+
+from embodichain.lab.devices.device import Device
+from embodichain.utils import logger
+
+if TYPE_CHECKING:
+    from embodichain.lab.sim.objects import Robot
+
+
+class DeviceController:
+    """Controller that bridges input devices (VR, keyboard, etc.) with robot control.
+
+    This controller is agnostic to the environment and can be used in both
+    gym environments and pure simulation contexts. It handles:
+    - Mapping device input to robot joint commands
+    - Joint limit enforcement
+    - Filtering and smoothing (if not handled by device)
+    - Multi-device support
+
+    Example:
+        # In gym environment
+        controller = DeviceController(robot=env.robot, device=vr_device)
+        action = controller.get_action()
+        env.step(action)
+
+        # In pure simulation
+        controller = DeviceController(robot=sim_robot, device=vr_device)
+        qpos = controller.get_action()
+        sim_robot.set_qpos(qpos)
+    """
+
+    def __init__(
+        self,
+        robot: Robot,
+        device: Optional[Device] = None,
+        device_name: str = "default",
+    ):
+        """Initialize Device Controller.
+
+        Args:
+            robot: Robot instance to control.
+            device: Input device (VR, keyboard, etc.). Can be None initially.
+            device_name: Name identifier for this device.
+        """
+        self.robot = robot
+        self._devices: Dict[str, Device] = {}
+        self._active_device_name: Optional[str] = None
+
+        if device is not None:
+            self.add_device(device, device_name, set_active=True)
+
+        # Joint mapping state: maps device joint name -> robot joint index
+        self._joint_mapping: Dict[str, int] = {}
+        self._mapping_initialized = False
+
+        logger.log_info(f"Device Controller initialized for robot: {robot.uid}")
+        logger.log_info(f"  Robot has {len(robot.joint_names)} joints")
+
+    def add_device(
+        self, device: Device, device_name: str, set_active: bool = False
+    ) -> None:
+        """Add a new input device.
+
+        Args:
+            device: Device instance to add.
+            device_name: Name identifier for the device.
+            set_active: Whether to set this as the active device.
+        """
+        self._devices[device_name] = device
+
+        if set_active or self._active_device_name is None:
+            self._active_device_name = device_name
+
+        logger.log_info(f"Added device '{device_name}' (Active: {set_active})")
+
+    def remove_device(self, device_name: str) -> None:
+        """Remove a device.
+
+        Args:
+            device_name: Name of the device to remove.
+        """
+        if device_name in self._devices:
+            del self._devices[device_name]
+            logger.log_info(f"Removed device '{device_name}'")
+
+            if self._active_device_name == device_name:
+                self._active_device_name = (
+                    list(self._devices.keys())[0] if self._devices else None
+                )
+
+    def set_active_device(self, device_name: str) -> None:
+        """Set the active input device.
+
+        Args:
+            device_name: Name of the device to activate.
+        """
+        if device_name not in self._devices:
+            logger.log_error(f"Device '{device_name}' not found")
+            return
+
+        self._active_device_name = device_name
+        logger.log_info(f"Active device set to '{device_name}'")
+
+    def get_action(
+        self, device_name: Optional[str] = None, as_dict: bool = False
+    ) -> Union[torch.Tensor, Dict[str, float], None]:
+        """Get robot action from device input.
+
+        Args:
+            device_name: Name of device to query. If None, uses active device.
+            as_dict: Whether to return action as dict (joint_name -> value).
+
+        Returns:
+            Robot action tensor (shape: [num_envs, num_joints]) or dict,
+            or None if no valid data available.
+        """
+        # Get device
+        device_name = device_name or self._active_device_name
+        if device_name is None or device_name not in self._devices:
+            return None
+
+        device = self._devices[device_name]
+
+        # Get device state
+        state = device.get_controller_state()
+        device_data = state.get("filtered_data") or state.get("raw_data")
+
+        if device_data is None:
+            return None
+
+        # Map device data to robot action
+        return self._map_device_to_robot(device_data, as_dict=as_dict)
+
+    def _map_device_to_robot(
+        self, device_data: Dict[str, float], as_dict: bool = False
+    ) -> Union[torch.Tensor, Dict[str, float], None]:
+        """Map device input to robot action.
+
+        Args:
+            device_data: Device joint data (joint_name -> value).
+            as_dict: Whether to return as dict instead of tensor.
+
+        Returns:
+            Robot action or None if mapping failed.
+        """
+        try:
+            robot_joint_names = self.robot.joint_names
+
+            # Build joint mapping on first call
+            if not self._mapping_initialized:
+                self._build_joint_mapping(device_data, robot_joint_names)
+
+            # Extract joint values based on mapping
+            joint_values = []
+            joint_indices = []
+            mapped_joints = {}
+
+            for device_joint, robot_idx in self._joint_mapping.items():
+                if device_joint in device_data:
+                    value = device_data[device_joint]
+                    joint_values.append(value)
+                    joint_indices.append(robot_idx)
+                    mapped_joints[robot_joint_names[robot_idx]] = value
+
+            if len(joint_indices) == 0:
+                return None
+
+            # Return as dict if requested
+            if as_dict:
+                return mapped_joints
+
+            # Convert to tensor and create full action
+            device_tensor = torch.tensor(
+                joint_values, dtype=torch.float32, device=self.robot.device
+            )
+            indices_tensor = torch.tensor(
+                joint_indices, dtype=torch.long, device=self.robot.device
+            )
+
+            # Get current robot qpos
+            current_qpos = self.robot.get_qpos()  # [num_envs, num_joints]
+
+            # Create action by updating controlled joints
+            action = current_qpos.clone()
+            action[:, indices_tensor] = device_tensor.unsqueeze(0)
+
+            # Enforce joint limits
+            action = self._enforce_joint_limits(action)
+
+            return action
+
+        except Exception as e:
+            logger.log_error(f"Error mapping device data to robot action: {e}")
+            return None
+
+    def _build_joint_mapping(
+        self, device_data: Dict[str, float], robot_joint_names: List[str]
+    ) -> None:
+        """Build mapping from device joints to robot joints.
+
+        Args:
+            device_data: Device joint data to determine available joints.
+            robot_joint_names: Robot joint names.
+        """
+        self._joint_mapping = {}
+
+        for robot_idx, robot_joint in enumerate(robot_joint_names):
+            # Direct name matching (can be extended with more sophisticated mapping)
+            if robot_joint in device_data:
+                self._joint_mapping[robot_joint] = robot_idx
+
+        self._mapping_initialized = True
+
+        logger.log_info(
+            f"Joint mapping initialized: {len(self._joint_mapping)} joints mapped"
+        )
+        logger.log_info(f"  Mapped joints: {list(self._joint_mapping.keys())}")
+
+    def _enforce_joint_limits(self, action: torch.Tensor) -> torch.Tensor:
+        """Enforce robot joint limits on action.
+
+        Args:
+            action: Action tensor [num_envs, num_joints].
+
+        Returns:
+            Clamped action tensor.
+        """
+        qpos_limits = self.robot.body_data.qpos_limits[0]  # [num_joints, 2]
+        return torch.clamp(action, qpos_limits[:, 0], qpos_limits[:, 1])
+
+    def reset(self) -> None:
+        """Reset controller state."""
+        # Reset all devices
+        for device in self._devices.values():
+            if hasattr(device, "reset"):
+                device.reset()
+
+        # Reset mapping
+        self._joint_mapping = {}
+        self._mapping_initialized = False
+
+        logger.log_info("Device Controller reset")
+
+    def get_device_info(self, device_name: Optional[str] = None) -> Dict[str, Any]:
+        """Get information about a device.
+
+        Args:
+            device_name: Name of device. If None, uses active device.
+
+        Returns:
+            Device information dict.
+        """
+        device_name = device_name or self._active_device_name
+        if device_name is None or device_name not in self._devices:
+            return {"error": "No device available"}
+
+        device = self._devices[device_name]
+        state = device.get_controller_state()
+
+        return {
+            "device_name": device_name,
+            "is_active": device_name == self._active_device_name,
+            "mapped_joints": list(self._joint_mapping.keys()),
+            "device_state": state,
+        }
+
+    def get_all_devices(self) -> List[str]:
+        """Get list of all registered device names."""
+        return list(self._devices.keys())
+
+    @property
+    def active_device(self) -> Optional[Device]:
+        """Get the currently active device."""
+        if self._active_device_name is None:
+            return None
+        return self._devices.get(self._active_device_name)

embodichain/lab/gym/envs/embodied_env.py

@@ -95,6 +95,19 @@ class EnvLightCfg:
     """Data pipeline configuration. Defaults to None.
     """
 
+    extensions: Union[Dict[str, Any], None] = None
+    """Extension parameters for task-specific configurations.
+    
+    This field can be used to pass additional parameters that are specific to certain environments
+    or tasks without modifying the base configuration class. For example:
+    - obs_mode: Observation mode (e.g., "state", "image")
+    - episode_length: Maximum episode length
+    - joint_limits: Joint limit constraints
+    - action_scale: Action scaling factor
+    - vr_joint_mapping: VR joint mapping for teleoperation
+    - control_frequency: Control frequency for VR teleoperation
+    """
+
     # Some helper attributes
     filter_visual_rand: bool = False
     """Whether to filter out visual randomization 
@@ -134,17 +147,9 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs):
 
         extensions = getattr(cfg, "extensions", {}) or {}
 
-        defaults = {
-            "obs_mode": "state",
-            "episode_length": 50,
-            "joint_limits": 0.5,
-            "action_scale": 0.1,
-        }
-
-        for name, default in defaults.items():
-            value = extensions.get(name, getattr(cfg, name, default))
+        for name, value in extensions.items():
             setattr(cfg, name, value)
-            setattr(self, name, getattr(cfg, name))
+            setattr(self, name, value)
 
         super().__init__(cfg, **kwargs)