hex-zmq-servers 0.3.9__py3-none-any.whl

hex_zmq_servers/mujoco/mujoco_base.py

@@ -0,0 +1,425 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-16
+################################################################
+
+import threading
+import numpy as np
+from collections import deque
+from abc import abstractmethod
+
+from ..device_base import HexDeviceBase
+from ..zmq_base import (
+    hex_zmq_ts_now,
+    HexZMQClientBase,
+    HexZMQServerBase,
+)
+
+NET_CONFIG = {
+    "ip": "127.0.0.1",
+    "port": 12345,
+    "realtime_mode": False,
+    "deque_maxlen": 10,
+    "client_timeout_ms": 200,
+    "server_timeout_ms": 1_000,
+    "server_num_workers": 4,
+}
+
+TAU = 2 * np.pi
+
+
+class HexMujocoBase(HexDeviceBase):
+
+    def __init__(self, realtime_mode: bool = False):
+        HexDeviceBase.__init__(self, realtime_mode)
+        self._dofs = None
+        self._limits = None
+        self._intri = None
+
+    def __del__(self):
+        HexDeviceBase.__del__(self)
+
+    def is_working(self) -> bool:
+        return self._working.is_set()
+
+    def get_dofs(self) -> np.ndarray:
+        self._wait_for_working()
+        return np.array(self._dofs, dtype=np.uint8)
+
+    def get_limits(self) -> np.ndarray:
+        self._wait_for_working()
+        return self._limits
+
+    def get_intri(self) -> np.ndarray:
+        self._wait_for_working()
+        return self._intri
+
+    @staticmethod
+    def _rads_normalize(rads: np.ndarray) -> np.ndarray:
+        return (rads + np.pi) % TAU - np.pi
+
+    @staticmethod
+    def _apply_pos_limits(
+        rads: np.ndarray,
+        lower_bound: np.ndarray,
+        upper_bound: np.ndarray,
+    ) -> np.ndarray:
+        normed_rads = HexMujocoBase._rads_normalize(rads)
+        outside = (normed_rads < lower_bound) | (normed_rads > upper_bound)
+        if not np.any(outside):
+            return normed_rads
+
+        lower_dist = np.fabs(
+            HexMujocoBase._rads_normalize(
+                (normed_rads - lower_bound)[outside]))
+        upper_dist = np.fabs(
+            HexMujocoBase._rads_normalize(
+                (normed_rads - upper_bound)[outside]))
+        choose_lower = lower_dist < upper_dist
+        choose_upper = ~choose_lower
+
+        outside_full = np.flatnonzero(outside)
+        outside_lower = outside_full[choose_lower]
+        outside_upper = outside_full[choose_upper]
+        normed_rads[outside_lower] = lower_bound[outside_lower]
+        normed_rads[outside_upper] = upper_bound[outside_upper]
+
+        return normed_rads
+
+    @abstractmethod
+    def work_loop(self, hex_queues: list[deque | threading.Event]):
+        raise NotImplementedError(
+            "`work_loop` should be implemented by the child class")
+
+    @abstractmethod
+    def close(self):
+        raise NotImplementedError(
+            "`close` should be implemented by the child class")
+
+
+class HexMujocoClientBase(HexZMQClientBase):
+
+    def __init__(self, net_config: dict = NET_CONFIG):
+        HexZMQClientBase.__init__(self, net_config)
+        self._states_seq = {
+            "robot": 0,
+            "obj": 0,
+        }
+        self._used_states_queue = {
+            "robot": 0,
+            "obj": 0,
+        }
+        self._states_queue = {
+            "robot": deque(maxlen=self._max_seq_num),
+            "obj": deque(maxlen=self._max_seq_num),
+        }
+        self._camera_seq = {
+            "rgb": 0,
+            "depth": 0,
+        }
+        self._used_camera_seq = {
+            "rgb": 0,
+            "depth": 0,
+        }
+        self._camera_queue = {
+            "rgb": deque(maxlen=self._deque_maxlen),
+            "depth": deque(maxlen=self._deque_maxlen),
+        }
+        self._cmds_seq = 0
+        self._cmds_queue = deque(maxlen=self._deque_maxlen)
+
+    def __del__(self):
+        HexZMQClientBase.__del__(self)
+
+    def reset(self):
+        self._states_seq = {
+            "robot": 0,
+            "obj": 0,
+        }
+        self._used_states_queue = {
+            "robot": 0,
+            "obj": 0,
+        }
+        self._camera_seq = {
+            "rgb": 0,
+            "depth": 0,
+        }
+        self._used_camera_seq = {
+            "rgb": 0,
+            "depth": 0,
+        }
+        self._cmds_seq = 0
+
+        hdr, _ = self.request({"cmd": "reset"})
+        return hdr
+
+    def seq_clear(self):
+        clear_hdr, _ = self.request({"cmd": "seq_clear"})
+        return clear_hdr
+
+    def get_dofs(self):
+        _, dofs = self.request({"cmd": "get_dofs"})
+        return dofs
+
+    def get_limits(self):
+        _, limits = self.request({"cmd": "get_limits"})
+        return limits
+
+    def get_states(self, robot_name: str | None = None, newest: bool = False):
+        try:
+            if self._realtime_mode or newest:
+                hdr, states = self._states_queue[robot_name][-1]
+                if self._used_states_queue[robot_name] != hdr["args"]:
+                    self._used_states_queue[robot_name] = hdr["args"]
+                    return hdr, states
+                else:
+                    return None, None
+            else:
+                return self._states_queue[robot_name].popleft()
+        except IndexError:
+            return None, None
+        except KeyError:
+            print(f"\033[91munknown robot name: {robot_name}\033[0m")
+            return None, None
+
+    def get_rgb(self, camera_name: str | None = None, newest: bool = False):
+        try:
+            if self._realtime_mode or newest:
+                hdr, img = self._camera_queue["rgb"][-1]
+                if self._used_camera_seq["rgb"] != hdr["args"]:
+                    self._used_camera_seq["rgb"] = hdr["args"]
+                    return hdr, img
+                else:
+                    return None, None
+            else:
+                return self._camera_queue["rgb"].popleft()
+        except IndexError:
+            return None, None
+
+    def get_depth(self, camera_name: str | None = None, newest: bool = False):
+        try:
+            if self._realtime_mode or newest:
+                hdr, img = self._camera_queue["depth"][-1]
+                if self._used_camera_seq["depth"] != hdr["args"]:
+                    self._used_camera_seq["depth"] = hdr["args"]
+                    return hdr, img
+                else:
+                    return None, None
+            else:
+                return self._camera_queue["depth"].popleft()
+        except IndexError:
+            return None, None
+
+    def set_cmds(self, cmds: np.ndarray):
+        self._cmds_queue.append(cmds)
+
+    def get_intri(self):
+        intri_hdr, intri = self.request({"cmd": "get_intri"})
+        return intri_hdr, intri
+
+    def _get_rgb_inner(self, camera_name: str | None = None):
+        return self._process_frame(camera_name, False)
+
+    def _get_depth_inner(self, camera_name: str | None = None):
+        return self._process_frame(camera_name, True)
+
+    def _process_frame(
+        self,
+        camera_name: str | None = None,
+        depth_flag: bool = False,
+    ):
+        req_cmd = f"get_{'depth' if depth_flag else 'rgb'}"
+        if camera_name is not None:
+            req_cmd += f"_{camera_name}"
+
+        hdr, img = self.request({
+            "cmd":
+            req_cmd,
+            "args": (1 + self._camera_seq['depth' if depth_flag else 'rgb']) %
+            self._max_seq_num,
+        })
+
+        try:
+            cmd = hdr["cmd"]
+            if cmd == f"{req_cmd}_ok":
+                if depth_flag:
+                    self._depth_seq = hdr["args"]
+                else:
+                    self._rgb_seq = hdr["args"]
+                return hdr, img
+            else:
+                return None, None
+        except KeyError:
+            print(f"\033[91m{hdr['cmd']} requires `cmd`\033[0m")
+            return None, None
+        except Exception as e:
+            print(f"\033[91m__process_frame failed: {e}\033[0m")
+            return None, None
+
+    def _get_states_inner(self, robot_name: str | None = None):
+        req_cmd = "get_states"
+        if robot_name is not None:
+            req_cmd += f"_{robot_name}"
+        hdr, states = self.request({
+            "cmd":
+            req_cmd,
+            "args": (1 + self._states_seq[robot_name]) % self._max_seq_num,
+        })
+        try:
+            if hdr["cmd"] == f"{req_cmd}_ok":
+                self._states_seq[robot_name] = hdr["args"]
+                return hdr, states
+            else:
+                return None, None
+        except KeyError:
+            print(f"\033[91m{hdr['cmd']} requires `cmd`\033[0m")
+            return None, None
+        except Exception as e:
+            print(f"\033[91m{req_cmd} failed: {e}\033[0m")
+            return None, None
+
+    def _set_cmds_inner(
+        self,
+        cmds: np.ndarray,
+    ) -> bool:
+        req_cmd = "set_cmds"
+        hdr, _ = self.request(
+            {
+                "cmd": req_cmd,
+                "ts": hex_zmq_ts_now(),
+                "args": self._cmds_seq,
+            },
+            cmds,
+        )
+        # print(f"{req_cmd} seq: {self._cmds_seq}")
+        try:
+            cmd = hdr["cmd"]
+            if cmd == f"{req_cmd}_ok":
+                self._cmds_seq = (self._cmds_seq + 1) % self._max_seq_num
+                return True
+            else:
+                return False
+        except KeyError:
+            print(f"\033[91m{hdr['cmd']} requires `cmd`\033[0m")
+            return False
+        except Exception as e:
+            print(f"\033[91m{req_cmd} failed: {e}\033[0m")
+            return False
+
+
+class HexMujocoServerBase(HexZMQServerBase):
+
+    def __init__(self, net_config: dict = NET_CONFIG):
+        HexZMQServerBase.__init__(self, net_config)
+        self._device: HexDeviceBase = None
+        self._states_queue = deque(maxlen=self._deque_maxlen)
+        self._obj_pose_queue = deque(maxlen=self._deque_maxlen)
+        self._cmds_queue = deque(maxlen=self._deque_maxlen)
+        self._cmds_seq = -1
+        self._rgb_queue = deque(maxlen=self._deque_maxlen)
+        self._depth_queue = deque(maxlen=self._deque_maxlen)
+        self._seq_clear_flag = False
+
+    def work_loop(self):
+        try:
+            self._device.work_loop([
+                self._states_queue,
+                self._obj_pose_queue,
+                self._cmds_queue,
+                self._rgb_queue,
+                self._depth_queue,
+                self._stop_event,
+            ])
+        finally:
+            self._device.close()
+
+    def _seq_clear(self):
+        self._seq_clear_flag = True
+        return True
+
+    def _get_states(self, recv_hdr: dict):
+        try:
+            seq = recv_hdr["args"]
+        except KeyError:
+            print(f"\033[91m{recv_hdr['cmd']} requires `args`\033[0m")
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+        try:
+            ts, count, states = self._states_queue[
+                -1] if self._realtime_mode else self._states_queue.popleft()
+        except IndexError:
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+        except Exception as e:
+            print(f"\033[91m{recv_hdr['cmd']} failed: {e}\033[0m")
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+        delta = (count - seq) % self._max_seq_num
+        if delta >= 0 and delta < 1e6:
+            return {
+                "cmd": f"{recv_hdr['cmd']}_ok",
+                "ts": ts,
+                "args": count
+            }, states
+        else:
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+    def _set_cmds(self, recv_hdr: dict, recv_buf: np.ndarray):
+        seq = recv_hdr.get("args", None)
+        if self._seq_clear_flag:
+            self._seq_clear_flag = False
+            self._cmds_seq = -1
+            return self.no_ts_hdr(recv_hdr, False), None
+
+        if seq is not None and seq > self._cmds_seq:
+            delta = (seq - self._cmds_seq) % self._max_seq_num
+            if delta >= 0 and delta < 1e6:
+                self._cmds_seq = seq
+                self._cmds_queue.append((recv_hdr["ts"], seq, recv_buf))
+                return self.no_ts_hdr(recv_hdr, True), None
+            else:
+                return self.no_ts_hdr(recv_hdr, False), None
+        else:
+            return self.no_ts_hdr(recv_hdr, False), None
+
+    def _get_frame(self, recv_hdr: dict):
+        try:
+            seq = recv_hdr["args"]
+        except KeyError:
+            print(f"\033[91m{recv_hdr['cmd']} requires `args`\033[0m")
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+        # get camera config
+        split_cmd = recv_hdr["cmd"].split("_")
+        depth_flag = split_cmd[1] == "depth"
+
+        try:
+            if depth_flag:
+                ts, count, img = self._depth_queue[
+                    -1] if self._realtime_mode else self._depth_queue.popleft(
+                    )
+            else:
+                ts, count, img = self._rgb_queue[
+                    -1] if self._realtime_mode else self._rgb_queue.popleft()
+        except IndexError:
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+        except Exception as e:
+            print(f"\033[91m{recv_hdr['cmd']} failed: {e}\033[0m")
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+        delta = (count - seq) % self._max_seq_num
+        if delta >= 0 and delta < 1e6:
+            return {
+                "cmd": f"{recv_hdr['cmd']}_ok",
+                "ts": ts,
+                "args": count
+            }, img
+        else:
+            return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
+
+    @abstractmethod
+    def _process_request(self, recv_hdr: dict, recv_buf: np.ndarray):
+        raise NotImplementedError(
+            "`_process_request` should be implemented by the child class")

hex_zmq_servers/robot/__init__.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-12
+################################################################
+
+from .robot_base import HexRobotBase, HexRobotClientBase, HexRobotServerBase
+from .dummy import HexRobotDummy, HexRobotDummyClient, HexRobotDummyServer
+from .gello import HexRobotGello, HexRobotGelloClient, HexRobotGelloServer
+from .hexarm import HexRobotHexarm, HexRobotHexarmClient, HexRobotHexarmServer, HEXARM_URDF_PATH_DICT
+
+__all__ = [
+    # path
+    "HEXARM_URDF_PATH_DICT",
+
+    # base
+    "HexRobotBase",
+    "HexRobotClientBase",
+    "HexRobotServerBase",
+
+    # dummy
+    "HexRobotDummy",
+    "HexRobotDummyClient",
+    "HexRobotDummyServer",
+
+    # gello
+    "HexRobotGello",
+    "HexRobotGelloClient",
+    "HexRobotGelloServer",
+
+    # hexarm
+    "HexRobotHexarm",
+    "HexRobotHexarmClient",
+    "HexRobotHexarmServer",
+]

hex_zmq_servers/robot/dummy/__init__.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-12
+################################################################
+
+from .robot_dummy import HexRobotDummy
+from .robot_dummy_cli import HexRobotDummyClient
+from .robot_dummy_srv import HexRobotDummyServer
+
+__all__ = [
+    "HexRobotDummy",
+    "HexRobotDummyClient",
+    "HexRobotDummyServer",
+]

hex_zmq_servers/robot/dummy/robot_dummy.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-12
+################################################################
+
+import threading
+import numpy as np
+from collections import deque
+
+from ..robot_base import HexRobotBase
+from ...zmq_base import (
+    hex_zmq_ts_now,
+    hex_zmq_ts_delta_ms,
+    HexRate,
+)
+from ...hex_launch import hex_log, HEX_LOG_LEVEL
+
+ROBOT_CONFIG = {
+    "dofs": [7],
+    "limits": [[[-1.0, 1.0]] * 3] * 7,
+    "states_init": [[0.0, 0.0, 0.0]] * 7,
+}
+
+
+class HexRobotDummy(HexRobotBase):
+
+    def __init__(
+        self,
+        robot_config: dict = ROBOT_CONFIG,
+        realtime_mode: bool = False,
+    ):
+        HexRobotBase.__init__(self, realtime_mode)
+
+        try:
+            self._dofs = robot_config["dofs"]
+            self._limits = np.array(robot_config["limits"])
+        except KeyError as ke:
+            missing_key = ke.args[0]
+            raise ValueError(
+                f"robot_config is not valid, missing key: {missing_key}")
+
+        # start work loop
+        self._working.set()
+
+    def __del__(self):
+        HexRobotBase.__del__(self)
+
+    def work_loop(self, hex_queues: list[deque | threading.Event]):
+        states_queue = hex_queues[0]
+        cmds_queue = hex_queues[1]
+        stop_event = hex_queues[2]
+
+        dummy_states = np.zeros((self._dofs[0], 3))
+        states_count = 0
+        last_cmds_seq = -1
+        rate = HexRate(1000)
+        while self._working.is_set() and not stop_event.is_set():
+            # states
+            states_queue.append((hex_zmq_ts_now(), states_count, dummy_states))
+            states_count = (states_count + 1) % self._max_seq_num
+
+            # cmds
+            cmds_pack = None
+            try:
+                cmds_pack = cmds_queue[
+                    -1] if self._realtime_mode else cmds_queue.popleft()
+            except IndexError:
+                pass
+            if cmds_pack is not None:
+                ts, seq, cmds = cmds_pack
+                if seq != last_cmds_seq:
+                    last_cmds_seq = seq
+                    if hex_zmq_ts_delta_ms(hex_zmq_ts_now(), ts) < 200.0:
+                        cmds = np.clip(
+                            cmds,
+                            self._limits[:, :, 0],
+                            self._limits[:, :, 1],
+                        )
+                        dummy_states = cmds.copy()
+
+            # sleep
+            rate.sleep()
+
+        # close
+        self.close()
+
+    def close(self):
+        if not self._working.is_set():
+            return
+        self._working.clear()
+        hex_log(HEX_LOG_LEVEL["info"], "HexRobotDummy closed")

hex_zmq_servers/robot/dummy/robot_dummy_cli.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-16
+################################################################
+
+from ..robot_base import HexRobotClientBase
+
+NET_CONFIG = {
+    "ip": "127.0.0.1",
+    "port": 12345,
+    "realtime_mode": False,
+    "deque_maxlen": 10,
+    "client_timeout_ms": 200,
+    "server_timeout_ms": 1_000,
+    "server_num_workers": 4,
+}
+
+
+class HexRobotDummyClient(HexRobotClientBase):
+
+    def __init__(
+        self,
+        net_config: dict = NET_CONFIG,
+    ):
+        HexRobotClientBase.__init__(self, net_config)
+        self._wait_for_working()

hex_zmq_servers/robot/dummy/robot_dummy_srv.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-12
+################################################################
+
+import numpy as np
+
+try:
+    from ..robot_base import HexRobotServerBase
+    from .robot_dummy import HexRobotDummy
+except (ImportError, ValueError):
+    import sys
+    from pathlib import Path
+    this_file = Path(__file__).resolve()
+    project_root = this_file.parents[3]
+    if str(project_root) not in sys.path:
+        sys.path.insert(0, str(project_root))
+    from hex_zmq_servers.robot.robot_base import HexRobotServerBase
+    from hex_zmq_servers.robot.dummy.robot_dummy import HexRobotDummy
+
+NET_CONFIG = {
+    "ip": "127.0.0.1",
+    "port": 12345,
+    "realtime_mode": False,
+    "deque_maxlen": 10,
+    "client_timeout_ms": 200,
+    "server_timeout_ms": 1_000,
+    "server_num_workers": 4,
+}
+
+ROBOT_CONFIG = {
+    "dofs": [7],
+    "limits": [[[-1.0, 1.0]] * 3] * 7,
+    "states_init": [[0.0, 0.0, 0.0]] * 7,
+}
+
+
+class HexRobotDummyServer(HexRobotServerBase):
+
+    def __init__(
+        self,
+        net_config: dict = NET_CONFIG,
+        params_config: dict = ROBOT_CONFIG,
+    ):
+        HexRobotServerBase.__init__(self, net_config)
+
+        # robot
+        self._device = HexRobotDummy(params_config,
+                                     net_config.get("realtime_mode", False))
+
+    def _process_request(self, recv_hdr: dict, recv_buf: np.ndarray):
+        if recv_hdr["cmd"] == "is_working":
+            return self.no_ts_hdr(recv_hdr, self._device.is_working()), None
+        elif recv_hdr["cmd"] == "seq_clear":
+            return self.no_ts_hdr(recv_hdr, self._seq_clear()), None
+        elif recv_hdr["cmd"] == "get_dofs":
+            dofs = self._device.get_dofs()
+            return self.no_ts_hdr(recv_hdr, dofs is not None), dofs
+        elif recv_hdr["cmd"] == "get_limits":
+            limits = self._device.get_limits()
+            return self.no_ts_hdr(recv_hdr, limits is not None), limits
+        elif recv_hdr["cmd"] == "get_states":
+            return self._get_states(recv_hdr)
+        elif recv_hdr["cmd"] == "set_cmds":
+            return self._set_cmds(recv_hdr, recv_buf)
+        else:
+            raise ValueError(f"unknown command: {recv_hdr['cmd']}")
+
+
+if __name__ == "__main__":
+    import argparse, json
+    from hex_zmq_servers.zmq_base import hex_server_helper
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cfg", type=str, required=True)
+    args = parser.parse_args()
+    cfg = json.loads(args.cfg)
+
+    hex_server_helper(cfg, HexRobotDummyServer)

hex_zmq_servers/robot/gello/__init__.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+################################################################
+# Copyright 2025 Dong Zhaorui. All rights reserved.
+# Author: Dong Zhaorui 847235539@qq.com
+# Date  : 2025-09-12
+################################################################
+
+from .robot_gello import HexRobotGello
+from .robot_gello_cli import HexRobotGelloClient
+from .robot_gello_srv import HexRobotGelloServer
+
+__all__ = [
+    "HexRobotGello",
+    "HexRobotGelloClient",
+    "HexRobotGelloServer",
+]