tairos-data-convert 1.0.3__py3-none-any.whl

data_pipeline/data_process/data_convert/data_convert.py

@@ -0,0 +1,142 @@
+import concurrent
+import os
+from pathlib import Path
+
+import yaml
+from loguru import logger
+import numpy as np
+
+from data_pipeline.data_process.utils.data_load import (
+    check_config_for_loading_dataset,
+    check_config_for_loading_data,
+    load_dataset,
+    load_data,
+)
+from data_pipeline.data_process.utils.data_read import BagReader
+from data_pipeline.data_process.utils.data_check import populate_dof_from_data
+from data_pipeline.data_process.utils.output_dataset import Hdf5Dataset
+from data_pipeline.data_process.utils.topic_mapping import get_data_for_value
+
+
+def check_config(config: dict):
+    if not config.get("dataset_id"):
+        raise ValueError("配置中未设置dataset_id")
+
+    if not config.get("fps"):
+        raise ValueError("配置中未设置fps")
+
+    output_dataset_dir = config.get("output_dataset_dir")
+    if not output_dataset_dir:
+        raise ValueError("配置中未设置output_dataset_dir")
+    os.makedirs(output_dataset_dir, exist_ok=True)
+
+    frame_structure = config.get("frame_structure")
+    if not frame_structure:
+        raise ValueError("配置中未设置frame_structure")
+
+    if not config.get("convert_num_workers"):
+        raise ValueError("配置中未设置convert_num_workers")
+
+    if not config.get("data_load"):
+        raise ValueError("配置中未设置data_load")
+    check_config_for_loading_data(config.get("data_load"))
+    check_config_for_loading_dataset(config.get("data_load"))
+
+
+def get_frame_data(frame_structure_values_set: set, frame_structure: dict, topic_to_msg: dict, data_source: int):
+    values_to_data = {}
+    for v in frame_structure_values_set:
+        values_to_data[v] = get_data_for_value(topic_to_msg, v, data_source)
+
+    frame_data = {}
+    for key, value in frame_structure.items():
+        frame_data[key] = np.concatenate([values_to_data[v] for v in value["values"]], axis=0)
+
+    return frame_data
+
+
+def convert(metadata, config: dict):
+    data_path = load_data(metadata, config.get("data_load"))
+
+    data_source = metadata.get("metadata", {}).get("data_source", 0)
+    frame_structure = config.get("frame_structure").copy()
+    try:
+        populate_dof_from_data(frame_structure, data_path, data_source)
+    except (KeyError, ValueError) as e:
+        return {
+            "success": False,
+            "error_message": str(e),
+        }
+
+    error_keys = []
+    for key, value in frame_structure.items():
+        if value["type"] == "float32" and value["dof"] == 0:
+            error_keys.append(key)
+    if error_keys:
+        return {
+            "success": False,
+            "error_message": f"frame_structure中存在dof为0的key: {error_keys}",
+        }
+
+    frame_structure_values_set = set()
+    for value in frame_structure.values():
+        for v in value["values"]:
+            frame_structure_values_set.add(v)
+
+    output_dataset_path = os.path.join(
+        config.get("output_dataset_dir"),
+        os.path.basename(data_path).rsplit('.', 1)[0] + ".hdf5"
+    )
+
+    with Hdf5Dataset(output_dataset_path, frame_structure) as output_dataset:
+        with BagReader(data_path) as bag_reader:
+            start_time = bag_reader.get_start_time()
+            end_time = bag_reader.get_end_time()
+            sample_timestamps = np.arange(start_time, end_time, 1 / config.get("fps"))
+
+            sample_index = 1
+            topic_to_msg = {}
+            for topic, msg, t in bag_reader.read_messages():
+                if t > sample_timestamps[sample_index]:
+                    try:
+                        frame_data = get_frame_data(
+                            frame_structure_values_set,
+                            frame_structure,
+                            topic_to_msg,
+                            data_source
+                        )
+                        output_dataset.add_frame(frame_data)
+                    except (KeyError, ValueError) as e:
+                        if t - start_time > 1:
+                            return {
+                                "success": False,
+                                "error_message": f"failed at relative time {t - start_time:2f}s: {str(e)}",
+                            }
+
+                    sample_index += 1
+                    if sample_index >= len(sample_timestamps):
+                        break
+
+                topic_to_msg[topic] = msg
+
+    return {
+        "success": True,
+    }
+
+
+def convert_data(config_path: str):
+    with open(config_path, 'r', encoding='utf-8') as f:
+        config = yaml.safe_load(f)
+    check_config(config)
+
+    metadatas = load_dataset(config.get("dataset_id"), config.get("data_load"))
+    # convert data in parallel
+    with concurrent.futures.ProcessPoolExecutor(max_workers=config.get("convert_num_workers")) as executor:
+        futures = {executor.submit(convert, metadata, config): metadata for metadata in metadatas}
+        for future in concurrent.futures.as_completed(futures):
+            result = future.result()
+            logger.info(f"converted {futures[future]['metadata']['name']}, result: {result}")
+
+
+if __name__ == "__main__":
+    convert_data(Path(__file__).parent / "config.yaml")

data_pipeline/data_process/utils/data_check.py

@@ -0,0 +1,40 @@
+from data_pipeline.data_process.utils.data_read import BagReader
+from data_pipeline.data_process.utils.topic_mapping import get_data_for_value
+
+
+def populate_dof(topic_to_msg, frame_structure, data_source):
+    for key, value in frame_structure.items():
+        dof = 0
+        for v in value["values"]:
+            data = get_data_for_value(topic_to_msg, v, data_source)
+            if value["type"] == "float32":
+                dof += data.shape[0]
+
+        if value["type"] == "float32":
+            frame_structure[key]["dof"] = dof
+    return
+
+
+def populate_dof_from_data(frame_structure, data_path, data_source):
+    with BagReader(data_path) as bag_reader:
+        start_time = bag_reader.get_start_time()
+        last_t = start_time
+        topic_to_msg = {}
+        last_error = ""
+        for topic, msg, t in bag_reader.read_messages():
+            topic_to_msg[topic] = msg
+
+            if t > start_time + 1:
+                raise ValueError(f"Cannot get all required values in the first second, last error: {last_error}")
+
+            if t > last_t + 0.1:
+                last_t = t
+                try:
+                    populate_dof(topic_to_msg, frame_structure, data_source)
+                except (KeyError, ValueError) as e:
+                    last_error = str(e)
+                    continue
+
+                return
+
+    raise ValueError(f"Cannot get all required values in the first second, last error: {last_error}")

data_pipeline/data_process/utils/data_load.py

@@ -0,0 +1,144 @@
+"""数据加载模块，支持通过goosefs或API获取数据集信息和bag路径"""
+import os
+import json
+import shutil
+from typing import List, Dict, Any, Tuple
+
+import requests
+
+from data_pipeline.py_api.api_utils import (
+    get_download_url_by_data_id,
+    get_all_data_from_dataset,
+)
+
+
+def check_config_for_loading_dataset(config: dict):
+    read_dataset_file = config.get("read_dataset_file")
+    if read_dataset_file is None:
+        raise ValueError("配置中未设置read_dataset_file")
+
+    if read_dataset_file and not config.get("goosefs_path"):
+        raise ValueError("read_dataset_file为true，但配置中未设置goosefs_path")
+
+
+def load_dataset(dataset_id: str, config: dict) -> List[Dict[str, Any]]:
+    """
+    根据配置选择从文件或API加载dataset
+
+    Args:
+        dataset_id: dataset ID
+
+    Returns:
+        metadata列表
+    """
+    if config.get("read_dataset_file"):
+        file_path = os.path.join(config.get("goosefs_path"), f"dataset_conf/{dataset_id}.json")
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Dataset文件不存在: {file_path}")
+
+        with open(file_path, "r", encoding='utf-8') as f:
+            return json.load(f)
+    else:
+        return get_all_data_from_dataset(dataset_id)
+
+
+def _get_file_key(metadata: Dict[str, Any]) -> str:
+    """从metadata中提取文件key"""
+    files = metadata.get("cos_storage", {}).get("files", [])
+    if len(files) == 0:
+        raise ValueError("failed to get cos_storage.files")
+
+    key = files[0].get("key", "")
+    if not key:
+        raise ValueError("failed to get cos_storage.files[0].key")
+
+    return key
+
+
+def _get_name_data_id_type(metadata: Dict[str, Any]) -> Tuple[str, str, str]:
+    key = _get_file_key(metadata)
+
+    metadata_name = metadata.get("metadata", {}).get("name")
+    if not metadata_name:
+        raise ValueError("failed to get metadata.name")
+
+    data_id = metadata.get("metadata", {}).get("data_id")
+    if not data_id:
+        raise ValueError("failed to get metadata.data_id")
+
+    return metadata_name, data_id, key.split(".")[-1]
+
+
+def get_path_on_goosefs_or_goosefsx(metadata: Dict[str, Any], config: dict) -> str:
+    key = _get_file_key(metadata)[1:]
+
+    goosefs_filepath = os.path.join(config.get("goosefs_path"), key)
+    if not os.path.exists(goosefs_filepath):
+        raise FileNotFoundError(f"goosefs文件不存在: {goosefs_filepath}")
+
+    if not config.get("use_goosefsx_as_cache"):
+        return goosefs_filepath
+
+    metadata_name, data_id, file_type = _get_name_data_id_type(metadata)
+    goosefsx_filepath = os.path.join(config.get("bag_cache_dir"), f"{metadata_name}_{data_id}.{file_type}")
+    shutil.copy(goosefs_filepath, goosefsx_filepath)
+    return goosefsx_filepath
+
+
+def download_bag_from_url(metadata: Dict[str, Any], config: dict) -> str:
+    metadata_name, data_id, file_type = _get_name_data_id_type(metadata)
+
+    download_url = get_download_url_by_data_id(data_id)
+    if not download_url:
+        raise ValueError(f"无法获取下载链接: {metadata_name}_{data_id}")
+
+    save_path = os.path.join(config.get("bag_download_dir"), f'{metadata_name}_{data_id}.{file_type}')
+    with requests.get(download_url, stream=True) as response:
+        response.raise_for_status()
+
+        with open(save_path, 'wb') as file:
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)
+
+    return save_path
+
+
+def check_config_for_loading_data(config: dict):
+    download_bag = config.get("download_bag")
+    if download_bag is None:
+        raise ValueError("配置中未设置download_bag")
+
+    bag_download_dir = config.get("bag_download_dir")
+    if download_bag and not bag_download_dir:
+        raise ValueError("download_bag为true，但配置中未设置bag_download_dir")
+    os.makedirs(bag_download_dir, exist_ok=True)
+
+    if not download_bag:
+        if not config.get("goosefs_path"):
+            raise ValueError("download_bag为false，但配置中未设置goosefs_path")
+
+        use_goosefsx_as_cache = config.get("use_goosefsx_as_cache")
+        if use_goosefsx_as_cache is None:
+            raise ValueError("download_bag为false，但配置中未设置use_goosefsx_as_cache")
+
+        bag_cache_dir = config.get("bag_cache_dir")
+        if use_goosefsx_as_cache and not bag_cache_dir:
+            raise ValueError("use_goosefsx_as_cache为true，但配置中未设置bag_cache_dir")
+        os.makedirs(bag_cache_dir, exist_ok=True)
+
+
+# Return path of the raw data
+def load_data(metadata: Dict[str, Any], config: dict) -> str:
+    """
+    根据配置选择下载或使用goosefs路径加载数据
+
+    Args:
+        metadata: metadata对象
+
+    Returns:
+        bag文件路径
+    """
+    if config.get("download_bag"):
+        return download_bag_from_url(metadata, config)
+    else:
+        return get_path_on_goosefs_or_goosefsx(metadata, config)

data_pipeline/data_process/utils/data_read.py

@@ -0,0 +1,90 @@
+"""数据读取模块，支持rosbag和mcap格式"""
+import os
+from typing import Iterator, Tuple, Optional, Any
+from pathlib import Path
+
+import rosbag
+from mcap_ros1.decoder import DecoderFactory as ROS1DecoderFactory
+from mcap.reader import make_reader
+
+
+class BagReader:
+    """数据读取器，支持rosbag和mcap格式"""
+
+    def __init__(self, bag_path: str):
+        """
+        初始化数据读取器
+
+        Args:
+            bag_path: bag或mcap文件路径
+        """
+        if not os.path.exists(bag_path):
+            raise FileNotFoundError(f"文件不存在: {bag_path}")
+
+        self.bag_path = bag_path
+        self.is_rosbag = bag_path.endswith(".bag")
+        self.is_mcap = bag_path.endswith(".mcap")
+
+        if not (self.is_rosbag or self.is_mcap):
+            raise ValueError(f"不支持的文件格式: {bag_path}")
+
+        if self.is_rosbag:
+            self.reader = rosbag.Bag(bag_path, 'r')
+        else:
+            self.file = open(bag_path, "rb")
+            self.reader = make_reader(self.file, decoder_factories=[ROS1DecoderFactory()])
+            self.summary = self.reader.get_summary()
+
+    def get_start_time(self) -> float:
+        """获取数据开始时间（秒）"""
+        if self.is_rosbag:
+            return self.reader.get_start_time()
+        else:
+            return self.summary.statistics.message_start_time / 1e9
+
+    def get_end_time(self) -> float:
+        """获取数据结束时间（秒）"""
+        if self.is_rosbag:
+            return self.reader.get_end_time()
+        else:
+            return self.summary.statistics.message_end_time / 1e9
+
+    def read_messages(self, topics: Optional[list] = None) -> Iterator[Tuple[str, Any, float]]:
+        """
+        读取消息
+
+        Args:
+            topics: 要读取的topic列表，如果为None则读取所有topic
+
+        Yields:
+            (topic, message, timestamp): topic名称、消息对象、时间戳（秒）
+        """
+        if self.is_rosbag:
+            for topic, msg, t in self.reader.read_messages(topics=topics):
+                yield topic, msg, t.to_sec()
+        else:
+            for _, channel, message, decoded_msg in self.reader.iter_decoded_messages(topics=topics):
+                yield channel.topic, decoded_msg, message.log_time / 1e9
+
+    def get_file_type(self) -> str:
+        """获取文件类型"""
+        if self.is_rosbag:
+            return "bag"
+        elif self.is_mcap:
+            return "mcap"
+        else:
+            return Path(self.bag_path).suffix[1:]  # 去掉点号
+
+    def close(self):
+        if self.is_rosbag:
+            self.reader.close()
+        else:
+            self.file.close()
+
+    def __enter__(self):
+        # open resources
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # close resources
+        self.close()

data_pipeline/data_process/utils/message_convert.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+
+def convert_message(msg, data_type=None):
+    msg_type = type(msg).__name__
+    if msg_type == "_sensor_msgs__CompressedImage":
+        return convert_compressed_image(msg)
+    if msg_type == "_data_msgs__ComponentObservation":
+        return convert_component_observation(msg, data_type)
+    if msg_type == "_data_msgs__ComponentAction":
+        return convert_component_action(msg, data_type)
+
+    raise ValueError(f"message type {msg_type} not supported")
+
+
+def convert_compressed_image(msg):
+    return np.frombuffer(msg.data, dtype=np.uint8)
+
+
+def convert_component_observation(msg, data_type):
+    if data_type == "pose":
+        pose = msg.multibody_pose.pose
+        return np.array([
+            pose.position.x, pose.position.y, pose.position.z,
+            pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w
+        ])
+    if data_type == "joint":
+        return np.array([state_msg.q for state_msg in msg.multibody_state.states])
+    raise ValueError(f"data_type {data_type} not supported")
+
+
+def convert_component_action(msg, data_type):
+    if data_type == "pose":
+        pose = msg.pose_command.pose
+        return np.array([
+            pose.position.x, pose.position.y, pose.position.z,
+            pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w
+        ])
+    if data_type == "joint":
+        return np.array(msg.joint_commands)
+    raise ValueError(f"data_type {data_type} not supported")

data_pipeline/data_process/utils/output_dataset.py

@@ -0,0 +1,56 @@
+import h5py
+import numpy as np
+
+
+class OutputDataset:
+    def __init__(self, path: str, frame_structure: dict):
+        self.path = path
+        self.frame_structure = frame_structure
+
+    def add_frame(self, frame_data: dict):
+        pass
+
+    def save_episode(self):
+        pass
+
+
+class Hdf5Dataset(OutputDataset):
+    def __init__(self, path: str, frame_structure: dict):
+        super().__init__(path, frame_structure)
+        self.hdf5_file = h5py.File(path, 'w')
+        self.datasets = {}
+
+        for key, value in self.frame_structure.items():
+            if value["type"] == "image":
+                self.datasets[key] = self.hdf5_file.create_dataset(
+                    key, (0,), maxshape=(None,), dtype=h5py.special_dtype(vlen=np.dtype('uint8')))
+            elif value["type"] == "string":
+                self.datasets[key] = self.hdf5_file.create_dataset(
+                    key, (0,), maxshape=(None,), dtype=h5py.string_dtype(encoding='utf-8'))
+            elif value["type"] == "float32":
+                self.datasets[key] = self.hdf5_file.create_dataset(
+                    key, (0, value["dof"]), maxshape=(None, value["dof"]), dtype='float32')
+            else:
+                raise ValueError(f"Unsupported frame structure type: {value['type']}")
+
+    def add_frame(self, frame_data: dict):
+        for key in self.frame_structure.keys():
+            self.datasets[key].resize(self.datasets[key].shape[0] + 1, axis=0)
+            self.datasets[key][-1] = frame_data[key]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.hdf5_file.__exit__(exc_type, exc_value, traceback)
+
+
+class LerobotDataset(OutputDataset):
+    def __init__(self, path: str):
+        super().__init__(path)
+
+    def add_frame(self, frame_data: dict):
+        pass
+
+    def save_episode(self):
+        pass

data_pipeline/data_process/utils/topic_mapping.py

@@ -0,0 +1,137 @@
+from data_pipeline.data_process.utils.message_convert import convert_message
+
+
+def get_data_for_value(topic_to_msg, value, data_source):
+    if value in topic_to_msg:
+        return convert_message(topic_to_msg[value])
+
+    if data_source == 1:
+        return get_umi_data_for_value(topic_to_msg, value)
+    elif data_source == 2:
+        return get_xtrainer_data_for_value(topic_to_msg, value)
+    elif data_source == 3:
+        return get_agibot_data_for_value(topic_to_msg, value)
+
+    raise ValueError(f"value {value} not supported for data source {data_source}")
+
+
+def get_umi_data_for_value(topic_to_msg, value):
+    if value == "head/camera/color":
+        return convert_message(topic_to_msg["/robot/data/head_realsense/color_image"])
+    if value == "left_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/left_hand_realsense/color_image"])
+    if value == "right_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/right_hand_realsense/color_image"])
+
+    if value == "left_gripper/state":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_left/observation"], data_type="joint")
+    if value == "right_gripper/state":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_right/observation"], data_type="joint")
+    if value == "left_gripper/command":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_left/action"], data_type="joint")
+    if value == "right_gripper/command":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_right/action"], data_type="joint")
+
+    if value == "left_wrist/pose/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/observation"], data_type="pose")
+    if value == "right_wrist/pose/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/observation"], data_type="pose")
+    if value == "left_wrist/pose/command":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/action"], data_type="pose")
+    if value == "right_wrist/pose/command":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/action"], data_type="pose")
+
+    if value == "left_wrist/joint/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/observation"], data_type="joint")
+    if value == "right_wrist/joint/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/observation"], data_type="joint")
+
+    raise ValueError(f"value {value} not supported for umi data")
+
+
+def get_agibot_data_for_value(topic_to_msg, value):
+    if value == "head/camera/color":
+        return convert_message(topic_to_msg["/robot/data/head_realsense/color_image"])
+    if value == "left_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/left_gripper/color_image"])
+    if value == "right_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/right_gripper/color_image"])
+
+    if value == "left_gripper/state":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_left/observation"], data_type="joint")
+    if value == "right_gripper/state":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_right/observation"], data_type="joint")
+    if value == "left_gripper/command":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_left/action"], data_type="joint")
+    if value == "right_gripper/command":
+        return convert_message(topic_to_msg["/robot/data/changingtek_hand_right/action"], data_type="joint")
+
+    if value == "left_wrist/pose/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/observation"], data_type="pose")
+    if value == "right_wrist/pose/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/observation"], data_type="pose")
+    if value == "left_wrist/pose/command":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/action"], data_type="pose")
+    if value == "right_wrist/pose/command":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/action"], data_type="pose")
+
+    if value == "left_wrist/joint/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_left/observation"], data_type="joint")
+    if value == "right_wrist/joint/state":
+        return convert_message(topic_to_msg["/robot/data/jaka_arm_right/observation"], data_type="joint")
+
+    raise ValueError(f"value {value} not supported for agibot data")
+
+
+def get_xtrainer_data_for_value(topic_to_msg, value):
+    if value == "head/camera/color":
+        return convert_message(topic_to_msg["/robot/data/head_realsense/color_image"])
+    if value == "left_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/left_hand_realsense/color_image"])
+    if value == "right_wrist/camera/color":
+        return convert_message(topic_to_msg["/robot/data/right_hand_realsense/color_image"])
+
+    if value == "left_gripper/state":
+        if "/robot/data/left_hand/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/left_hand/observation"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_left_hand/observation"], data_type="joint")
+    if value == "right_gripper/state":
+        if "/robot/data/right_hand/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/right_hand/observation"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_right_hand/observation"], data_type="joint")
+    if value == "left_gripper/command":
+        if "/robot/data/left_hand/action" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/left_hand/action"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_left_hand/action"], data_type="joint")
+    if value == "right_gripper/command":
+        if "/robot/data/right_hand/action" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/right_hand/action"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_right_hand/action"], data_type="joint")
+
+    if value == "left_wrist/pose/state":
+        if "/robot/data/left_arm/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/left_arm/observation"], data_type="pose")
+        return convert_message(topic_to_msg["/robot/data/aloha_left_arm/observation"], data_type="pose")
+    if value == "right_wrist/pose/state":
+        if "/robot/data/right_arm/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/right_arm/observation"], data_type="pose")
+        return convert_message(topic_to_msg["/robot/data/aloha_right_arm/observation"], data_type="pose")
+
+    if value == "left_wrist/joint/state":
+        if "/robot/data/left_arm/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/left_arm/observation"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_left_arm/observation"], data_type="joint")
+    if value == "right_wrist/joint/state":
+        if "/robot/data/right_arm/observation" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/right_arm/observation"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_right_arm/observation"], data_type="joint")
+    if value == "left_wrist/joint/command":
+        if "/robot/data/left_arm/action" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/left_arm/action"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_left_arm/action"], data_type="joint")
+    if value == "right_wrist/joint/command":
+        if "/robot/data/right_arm/action" in topic_to_msg:
+            return convert_message(topic_to_msg["/robot/data/right_arm/action"], data_type="joint")
+        return convert_message(topic_to_msg["/robot/data/aloha_right_arm/action"], data_type="joint")
+
+    raise ValueError(f"value {value} not supported for xtrainer data")