"""MuJoCo Sensors Module.
This module provides force-torque sensor implementations for the MuJoCo simulator.
It mirrors the base sensor interfaces defined in `robotblockset.sensors` and exposes
MuJoCo-backed sensor data through the same API.
Copyright (c) 2024 Jozef Stefan Institute
Authors: Leon Zlajpah.
"""
import numpy as np
from typing import Any, Optional
from robotblockset.sensors import force_torque_sensor
from robotblockset.mujoco.mujoco_api import mjInterface
[docs]
class ft_sensor(force_torque_sensor):
"""MuJoCo-backed force-torque sensor interface using the socket-based server API."""
[docs]
def __init__(self, scene: Optional[mjInterface] = None, host: str = "localhost", sensor_force_name: Optional[str] = None, sensor_torque_name: Optional[str] = None, **kwargs: Any) -> None:
"""
Initialize a MuJoCo force-torque sensor.
Parameters
----------
scene : mjInterface, optional
Existing MuJoCo interface instance. If None, a new connection is created.
host : str, optional
Hostname of the MuJoCo simulator.
sensor_force_name : str, optional
Name of the force sensor in MuJoCo. The default is ``"force"``.
sensor_torque_name : str, optional
Name of the torque sensor in MuJoCo. The default is ``"torque"``.
**kwargs : dict
Additional keyword arguments for the base force-torque sensor class.
"""
force_torque_sensor.__init__(self, **kwargs)
self.Name = "FTSensor_MuJoCo"
if sensor_force_name is None:
sensor_force_name = "force"
if sensor_torque_name is None:
sensor_torque_name = "torque"
if not isinstance(sensor_force_name, str):
raise TypeError("sensor_force_name must be a string")
if not isinstance(sensor_torque_name, str):
raise TypeError("sensor_torque_name must be a string")
self._owns_scene = scene is None
if scene is None:
self.scene = mjInterface(host=host)
self._connected = False
else:
self.scene = scene
if self.scene.mj_connected() == 0:
if self.scene.mj_connect() == 0:
self._connected = True
else:
raise RuntimeError("Connection to MuJoCo simulator failed")
else:
self._connected = True
self.DebugMessage("FT sensor connected to MuJoCo")
self.tsamp = 0.01
self.SensorForceName = sensor_force_name
self.SensorTorqueName = sensor_torque_name
self._info = None
self._sensor_layout_signature = None
self._SensorHandles: Optional[list[int]] = None
self.Init()
[docs]
def Init(self) -> None:
"""
Initialize MuJoCo sensor handles.
Returns
-------
None
"""
self._initialize_handles(self.scene.mj_info())
self.GetState()
self.DebugMessage("Initialized")
@staticmethod
def _get_sensor_layout_signature(info: Any, sensor_ids: tuple[int, int]) -> tuple[int, int, tuple[int, ...], tuple[int, ...], tuple[int, int]]:
"""Return a signature describing the sensor layout of a MuJoCo model."""
return (
int(info.nsensor),
int(info.nsensordata),
tuple(int(value) for value in info.sensor_adr),
tuple(int(value) for value in info.sensor_dim),
sensor_ids,
)
def _get_sensor_ids(self) -> tuple[int, int]:
"""Resolve the configured force and torque sensor names."""
return (
self.scene.mj_name2id("sensor", self.SensorForceName),
self.scene.mj_name2id("sensor", self.SensorTorqueName),
)
def _initialize_handles(self, info: Any, sensor_ids: Optional[tuple[int, int]] = None) -> None:
"""Resolve and validate the force and torque sensor handles."""
if sensor_ids is None:
sensor_ids = self._get_sensor_ids()
self._info = info
self._sensor_layout_signature = self._get_sensor_layout_signature(info, sensor_ids)
self._SensorHandles = None
if any(sensor_id < 0 for sensor_id in sensor_ids):
return
handles: list[int] = []
for sensor_id, sensor_name, quantity in zip(
sensor_ids,
(self.SensorForceName, self.SensorTorqueName),
("force", "torque"),
):
if sensor_id >= info.nsensor:
raise RuntimeError(f"Invalid MuJoCo sensor ID {sensor_id} for '{sensor_name}'")
address = int(info.sensor_adr[sensor_id])
dimension = int(info.sensor_dim[sensor_id])
if dimension != 3:
raise ValueError(f"MuJoCo {quantity} sensor '{sensor_name}' must have dimension 3, got {dimension}")
if address < 0 or address + dimension > info.nsensordata:
raise RuntimeError(f"Invalid MuJoCo data range for sensor '{sensor_name}'")
handles.extend(range(address, address + dimension))
self._SensorHandles = handles
[docs]
def GetRawFT(self) -> np.ndarray:
"""
Read raw force-torque data from MuJoCo sensors.
Returns
-------
np.ndarray
Force-torque vector (6,) or NaNs if unavailable.
"""
info = self.scene.mj_info()
sensor_ids = self._get_sensor_ids()
if self._sensor_layout_signature != self._get_sensor_layout_signature(info, sensor_ids):
self._initialize_handles(info, sensor_ids)
if info.nsensor > 0 and self._SensorHandles is not None:
sensor = self.scene.mj_get_sensor()
sensor_data = np.asarray(sensor.sensordata).reshape(-1)
if int(sensor.nsensordata) != sensor_data.size:
raise RuntimeError("MuJoCo sensor-data size does not match its descriptor")
if self._SensorHandles and max(self._SensorHandles) >= sensor_data.size:
raise RuntimeError("MuJoCo sensor handles exceed the returned sensor-data array")
self.SensorData = np.take(sensor_data, self._SensorHandles).astype(float, copy=True)
else:
self.SensorData = np.full(6, np.nan)
return self.SensorData.copy()
[docs]
def Close(self) -> None:
"""Close the MuJoCo connection if it was created by this sensor."""
if getattr(self, "_owns_scene", False) and getattr(self, "scene", None) is not None:
self.scene.mj_close()
self._connected = False
def __del__(self) -> None:
"""Release an owned MuJoCo connection and detach the sensor."""
try:
self.Close()
except Exception:
pass
try:
force_torque_sensor.__del__(self)
except Exception:
pass