"""PyMuJoCo Sensors Module.
This module provides force-torque sensor implementations for the Python MuJoCo simulator.
It mirrors the base sensor interfaces defined in `robotblockset.sensors` and exposes
PyMuJoCo-backed sensor data through the same API.
Copyright (c) 2025- 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.scene_pymujoco import mujoco_scene
[docs]
class ft_sensor(force_torque_sensor):
"""PyMuJoCo-backed force-torque sensor interface operating on a :class:`~robotblockset.mujoco.scene_pymujoco.mujoco_scene`."""
[docs]
def __init__(self, sensor_name: str = "FTSensor", scene: Optional[mujoco_scene] = None, sensor_force_name: Optional[str] = None, sensor_torque_name: Optional[str] = None, **kwargs: Any) -> None:
"""
Initialize a PyMuJoCo force-torque sensor.
Parameters
----------
sensor_name : str, optional
Sensor name prefix used for this instance.
scene : :class:`~robotblockset.mujoco.scene_pymujoco.mujoco_scene`, optional
MuJoCo scene instance to read sensors from. Must be provided.
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 = sensor_name + ":PyMuJoCo"
if scene is None:
raise ValueError("MuJoCo scene is not defined")
self.scene = scene
self._connected = True
self.DebugMessage(f"Sensor {sensor_name} connected to PyMuJoCo")
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.tsamp = 0.01
self.SensorForceName = sensor_force_name
self.SensorTorqueName = sensor_torque_name
self._SensorNamesList: list[str] = []
self._SensorHandles: Optional[list[int]] = None
self._model_signature: Optional[tuple[int, tuple[str, ...], tuple[int, ...], tuple[int, ...]]] = None
self.Init()
[docs]
def Init(self) -> None:
"""
Initialize MuJoCo sensor handles.
Returns
-------
None
"""
with self.scene.lock():
self._initialize_handles_locked()
self.GetState()
self.DebugMessage("Initialized")
def _get_model_signature_locked(self) -> tuple[int, tuple[str, ...], tuple[int, ...], tuple[int, ...]]:
"""Return the sensor-layout signature of the current MuJoCo model."""
model = self.scene.model
names = tuple(model.sensor(i).name for i in range(model.nsensor))
addresses = tuple(int(x) for x in np.asarray(model.sensor_adr).reshape(-1))
dimensions = tuple(int(x) for x in np.asarray(model.sensor_dim).reshape(-1))
return id(model), names, addresses, dimensions
def _initialize_handles_locked(self) -> None:
"""Resolve and validate force and torque handles while the scene is locked."""
model = self.scene.model
self._model_signature = self._get_model_signature_locked()
self._SensorNamesList = list(self._model_signature[1])
self._SensorHandles = None
if self.SensorForceName not in self._SensorNamesList or self.SensorTorqueName not in self._SensorNamesList:
return
handles: list[int] = []
for sensor_name, quantity in ((self.SensorForceName, "force"), (self.SensorTorqueName, "torque")):
sensor_id = int(model.sensor(sensor_name).id)
address = int(model.sensor_adr[sensor_id])
dimension = int(model.sensor_dim[sensor_id])
if dimension != 3:
raise ValueError(f"MuJoCo {quantity} sensor '{sensor_name}' must have dimension 3, got {dimension}")
handles.extend(range(address, address + dimension))
self._SensorHandles = handles
[docs]
def GetRawFT(self) -> np.ndarray:
"""
Read raw force-torque data from PyMuJoCo sensors.
Returns
-------
np.ndarray
Force-torque vector (6,) or NaNs if unavailable.
"""
with self.scene.lock():
if self._model_signature != self._get_model_signature_locked():
self._initialize_handles_locked()
if self._SensorHandles is None:
self.SensorData = np.full(6, np.nan)
else:
self.SensorData = np.take(self.scene.data.sensordata, self._SensorHandles).astype(float, copy=True)
return self.SensorData.copy()