space-mouse-在diffusion_policy项目中替换原有的代码

主要问题是 diffusion_policy 依赖的spnav比较挑 python版本，我们进行替换 ； 具体可以参考文档 space-mouse-在linux上初始化安装和测试过程

import pyspacemouse
from threading import Thread, Event
import numpy as np
import time

# 拷贝参考自 https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/real_world/spacemouse.py
# 删除了老的 spnav， 换成新的库 pyspacemouse，避免兼容性问题
class Spacemouse(Thread):
    def __init__(self, max_value=500, deadzone=(0, 0, 0, 0, 0, 0), dtype=np.float32):
        """
        Continuously listen to 3D connection space naviagtor events
        and update the latest state.

        max_value: {300, 500} 300 for wired version and 500 for wireless
        deadzone: [0,1], number or tuple, axis with value lower than this value will stay at 0

        front
        z
        ^   _
        |  (O) space mouse
        |
        *----->x right
        y
        """
        if np.issubdtype(type(deadzone), np.number):
            deadzone = np.full(6, fill_value=deadzone, dtype=dtype)
        else:
            deadzone = np.array(deadzone, dtype=dtype)
        assert (deadzone >= 0).all()

        super().__init__()
        self.stop_event = Event()
        self.max_value = max_value
        self.dtype = dtype
        self.deadzone = deadzone
        self.motion_state = np.zeros(6, dtype=dtype)
        self.button_state = {0: False, 1: False}
        self.tx_zup_spnav = np.array([
            [0, 0, -1],
            [1, 0, 0],
            [0, 1, 0]
        ], dtype=dtype)

    def get_motion_state(self):
        state = np.array(self.motion_state, dtype=self.dtype) / self.max_value
        is_dead = (-self.deadzone < state) & (state < self.deadzone)
        state[is_dead] = 0
        return state

    def get_motion_state_transformed(self):
        """
        Return in right-handed coordinate
        z
        *------>y right
        |   _
        |  (O) space mouse
        v
        x
        back

        """
        state = self.get_motion_state()
        tf_state = np.zeros_like(state)
        tf_state[:3] = self.tx_zup_spnav @ state[:3]
        tf_state[3:] = self.tx_zup_spnav @ state[3:]
        return tf_state

    def is_button_pressed(self, button_id):
        return self.button_state.get(button_id, False)

    def stop(self):
        self.stop_event.set()
        self.join()

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()

    def run(self):
        pyspacemouse.open()
        try:
            while not self.stop_event.is_set():
                state = pyspacemouse.read()
                if state is not None:
                    self.motion_state[0] = int(state.x * self.max_value)
                    self.motion_state[1] = int(state.y * self.max_value)
                    self.motion_state[2] = int(state.z * self.max_value)
                    
                    self.motion_state[3] = int(state.pitch * self.max_value)
                    self.motion_state[4] = int(state.yaw * self.max_value)
                    self.motion_state[5] = int(state.roll * self.max_value)
                    
                    if state.buttons:
                        for i, pressed in enumerate(state.buttons):
                            self.button_state[i] = bool(pressed)
                time.sleep(1 / 200)
        finally:
            pyspacemouse.close()


def test():
    with Spacemouse(deadzone=0.3) as sm:
        for i in range(2000):
            print(sm.get_motion_state_transformed())
            print(sm.is_button_pressed(0))
            time.sleep(1 / 100)


if __name__ == '__main__':
    test()