hex-zmq-servers 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

hex_zmq_servers/mujoco/e3_desktop/mujoco_e3_desktop_cli.py

@@ -7,7 +7,8 @@
 ################################################################
 
 import numpy as np
-from ...zmq_base import hex_zmq_ts_now
+from collections import deque
+from ...zmq_base import HexRate, hex_zmq_ts_now
 from ..mujoco_base import HexMujocoClientBase
 
 NET_CONFIG = {
@@ -17,6 +18,15 @@ NET_CONFIG = {
     "server_timeout_ms": 1_000,
     "server_num_workers": 4,
 }
+RECV_CONFIG = {
+    "head_rgb": True,
+    "head_depth": True,
+    "left_rgb": True,
+    "left_depth": True,
+    "right_rgb": True,
+    "right_depth": True,
+    "obj": True,
+}
 
 
 class HexMujocoE3DesktopClient(HexMujocoClientBase):
@@ -24,8 +34,10 @@ class HexMujocoE3DesktopClient(HexMujocoClientBase):
     def __init__(
         self,
         net_config: dict = NET_CONFIG,
+        recv_config: dict = RECV_CONFIG,
     ):
         HexMujocoClientBase.__init__(self, net_config)
+        self.__recv_config = recv_config
         self._cmds_seq = {
             "left": 0,
             "right": 0,
@@ -38,38 +50,24 @@ class HexMujocoE3DesktopClient(HexMujocoClientBase):
             "right_rgb": 0,
             "right_depth": 0,
         }
-
-    def set_cmds(
-        self,
-        cmds: np.ndarray,
-        robot_name: str | None = None,
-    ) -> bool:
-        req_cmd = "set_cmds"
-        if robot_name is not None:
-            req_cmd += f"_{robot_name}"
-        hdr, _ = self.request(
-            {
-                "cmd": req_cmd,
-                "ts": hex_zmq_ts_now(),
-                "args": self._cmds_seq[robot_name],
-            },
-            cmds,
-        )
-        # print(f"{req_cmd} seq: {self._cmds_seq[robot_name]}")
-        try:
-            cmd = hdr["cmd"]
-            if cmd == f"{req_cmd}_ok":
-                self._cmds_seq[robot_name] = (self._cmds_seq[robot_name] +
-                                              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
+        self._states_queue = {
+            "left": deque(maxlen=10),
+            "right": deque(maxlen=10),
+            "obj": deque(maxlen=10),
+        }
+        self._camera_queue = {
+            "head_rgb": deque(maxlen=10),
+            "head_depth": deque(maxlen=10),
+            "left_rgb": deque(maxlen=10),
+            "left_depth": deque(maxlen=10),
+            "right_rgb": deque(maxlen=10),
+            "right_depth": deque(maxlen=10),
+        }
+        self._cmds_queue = {
+            "left": deque(maxlen=10),
+            "right": deque(maxlen=10),
+        }
+        self._wait_for_working()
 
     def reset(self):
         HexMujocoClientBase.reset(self)
@@ -86,6 +84,35 @@ class HexMujocoE3DesktopClient(HexMujocoClientBase):
             "right_depth": 0,
         }
 
+    def get_rgb(self, camera_name: str | None = None, newest: bool = False):
+        name = f"{camera_name}_rgb"
+        try:
+            return self._camera_queue[name].popleft(
+            ) if not newest else self._camera_queue[name][-1]
+        except IndexError:
+            return None, None
+        except KeyError:
+            print(f"\033[91munknown camera name: {name}\033[0m")
+            return None, None
+
+    def get_depth(self, camera_name: str | None = None, newest: bool = False):
+        name = f"{camera_name}_depth"
+        try:
+            return self._camera_queue[name].popleft(
+            ) if not newest else self._camera_queue[name][-1]
+        except IndexError:
+            return None, None
+        except KeyError:
+            print(f"\033[91munknown camera name: {name}\033[0m")
+            return None, None
+
+    def set_cmds(
+        self,
+        cmds: np.ndarray,
+        robot_name: str | None = None,
+    ):
+        self._cmds_queue[robot_name].append(cmds)
+
     def _process_frame(
         self,
         camera_name: str | None = None,
@@ -116,3 +143,91 @@ class HexMujocoE3DesktopClient(HexMujocoClientBase):
         except Exception as e:
             print(f"\033[91m__process_frame failed: {e}\033[0m")
             return None, None
+
+    def _set_cmds_inner(
+        self,
+        cmds: np.ndarray,
+        robot_name: str | None = None,
+    ) -> bool:
+        req_cmd = "set_cmds"
+        if robot_name is not None:
+            req_cmd += f"_{robot_name}"
+        hdr, _ = self.request(
+            {
+                "cmd": req_cmd,
+                "ts": hex_zmq_ts_now(),
+                "args": self._cmds_seq[robot_name],
+            },
+            cmds,
+        )
+        # print(f"{req_cmd} seq: {self._cmds_seq[robot_name]}")
+        try:
+            cmd = hdr["cmd"]
+            if cmd == f"{req_cmd}_ok":
+                self._cmds_seq[robot_name] = (self._cmds_seq[robot_name] +
+                                              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
+
+    def _recv_loop(self):
+        rate = HexRate(2000)
+        image_trig_cnt = 0
+        while self._recv_flag:
+            hdr, states = self._get_states_inner("left")
+            if hdr is not None:
+                self._states_queue["left"].append((hdr, states))
+            hdr, states = self._get_states_inner("right")
+            if hdr is not None:
+                self._states_queue["right"].append((hdr, states))
+            if self.__recv_config["obj"]:
+                hdr, obj_pose = self._get_states_inner("obj")
+                if hdr is not None:
+                    self._states_queue["obj"].append((hdr, obj_pose))
+
+            image_trig_cnt += 1
+            if image_trig_cnt >= 10:
+                image_trig_cnt = 0
+                if self.__recv_config["head_rgb"]:
+                    hdr, img = self._get_rgb_inner("head")
+                    if hdr is not None:
+                        self._camera_queue["head_rgb"].append((hdr, img))
+                if self.__recv_config["head_depth"]:
+                    hdr, img = self._get_depth_inner("head")
+                    if hdr is not None:
+                        self._camera_queue["head_depth"].append((hdr, img))
+                if self.__recv_config["left_rgb"]:
+                    hdr, img = self._get_rgb_inner("left")
+                    if hdr is not None:
+                        self._camera_queue["left_rgb"].append((hdr, img))
+                if self.__recv_config["left_depth"]:
+                    hdr, img = self._get_depth_inner("left")
+                    if hdr is not None:
+                        self._camera_queue["left_depth"].append((hdr, img))
+                if self.__recv_config["right_rgb"]:
+                    hdr, img = self._get_rgb_inner("right")
+                    if hdr is not None:
+                        self._camera_queue["right_rgb"].append((hdr, img))
+                if self.__recv_config["right_depth"]:
+                    hdr, img = self._get_depth_inner("right")
+                    if hdr is not None:
+                        self._camera_queue["right_depth"].append((hdr, img))
+
+            try:
+                cmds = self._cmds_queue["left"].popleft()
+                _ = self._set_cmds_inner(cmds, "left")
+            except IndexError:
+                pass
+            try:
+                cmds = self._cmds_queue["right"].popleft()
+                _ = self._set_cmds_inner(cmds, "right")
+            except IndexError:
+                pass
+
+            rate.sleep()

hex_zmq_servers/mujoco/e3_desktop/mujoco_e3_desktop_srv.py

@@ -7,7 +7,7 @@
 ################################################################
 
 import numpy as np
-from hex_zmq_servers.zmq_base import HexSafeValue
+from collections import deque
 
 try:
     from ..mujoco_base import HexMujocoServerBase
@@ -31,7 +31,7 @@ NET_CONFIG = {
 }
 
 MUJOCO_CONFIG = {
-    "states_rate": 500,
+    "states_rate": 1000,
     "img_rate": 30,
     "headless": False,
     "sens_ts": True,
@@ -53,32 +53,32 @@ class HexMujocoE3DesktopServer(HexMujocoServerBase):
         # values
         self._cmds_left_seq = -1
         self._cmds_right_seq = -1
-        self._states_left_value = HexSafeValue()
-        self._states_right_value = HexSafeValue()
-        self._states_obj_value = HexSafeValue()
-        self._cmds_left_value = HexSafeValue()
-        self._cmds_right_value = HexSafeValue()
-        self._rgb_head_value = HexSafeValue()
-        self._depth_head_value = HexSafeValue()
-        self._rgb_left_value = HexSafeValue()
-        self._depth_left_value = HexSafeValue()
-        self._rgb_right_value = HexSafeValue()
-        self._depth_right_value = HexSafeValue()
+        self._states_left_queue = deque(maxlen=10)
+        self._states_right_queue = deque(maxlen=10)
+        self._states_obj_queue = deque(maxlen=10)
+        self._cmds_left_queue = deque(maxlen=10)
+        self._cmds_right_queue = deque(maxlen=10)
+        self._rgb_head_queue = deque(maxlen=10)
+        self._depth_head_queue = deque(maxlen=10)
+        self._rgb_left_queue = deque(maxlen=10)
+        self._depth_left_queue = deque(maxlen=10)
+        self._rgb_right_queue = deque(maxlen=10)
+        self._depth_right_queue = deque(maxlen=10)
 
     def work_loop(self):
         try:
             self._device.work_loop([
-                self._states_left_value,
-                self._states_right_value,
-                self._states_obj_value,
-                self._cmds_left_value,
-                self._cmds_right_value,
-                self._rgb_head_value,
-                self._depth_head_value,
-                self._rgb_left_value,
-                self._depth_left_value,
-                self._rgb_right_value,
-                self._depth_right_value,
+                self._states_left_queue,
+                self._states_right_queue,
+                self._states_obj_queue,
+                self._cmds_left_queue,
+                self._cmds_right_queue,
+                self._rgb_head_queue,
+                self._depth_head_queue,
+                self._rgb_left_queue,
+                self._depth_left_queue,
+                self._rgb_right_queue,
+                self._depth_right_queue,
                 self._stop_event,
             ])
         finally:
@@ -94,16 +94,18 @@ class HexMujocoE3DesktopServer(HexMujocoServerBase):
         # get robot name
         robot_name = recv_hdr["cmd"].split("_")[2]
         if robot_name == "left":
-            value = self._states_left_value
+            queue = self._states_left_queue
         elif robot_name == "right":
-            value = self._states_right_value
+            queue = self._states_right_queue
         elif robot_name == "obj":
-            value = self._states_obj_value
+            queue = self._states_obj_queue
         else:
             raise ValueError(f"unknown robot name: {robot_name}")
 
         try:
-            ts, count, states = value.get()
+            ts, count, 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
@@ -129,10 +131,10 @@ class HexMujocoE3DesktopServer(HexMujocoServerBase):
         # get robot name
         robot_name = recv_hdr["cmd"].split("_")[2]
         if robot_name == "left":
-            value = self._cmds_left_value
+            queue = self._cmds_left_queue
             cmds_seq = self._cmds_left_seq
         elif robot_name == "right":
-            value = self._cmds_right_value
+            queue = self._cmds_right_queue
             cmds_seq = self._cmds_right_seq
         else:
             raise ValueError(f"unknown robot name: {robot_name}")
@@ -144,7 +146,7 @@ class HexMujocoE3DesktopServer(HexMujocoServerBase):
                     self._cmds_left_seq = cmds_seq
                 elif robot_name == "right":
                     self._cmds_right_seq = cmds_seq
-                value.set((recv_hdr["ts"], seq, recv_buf))
+                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
@@ -163,16 +165,18 @@ class HexMujocoE3DesktopServer(HexMujocoServerBase):
         depth_flag = split_cmd[1] == "depth"
         camera_name = split_cmd[2]
         if camera_name == "head":
-            value = self._rgb_head_value if not depth_flag else self._depth_head_value
+            queue = self._rgb_head_queue if not depth_flag else self._depth_head_queue
         elif camera_name == "left":
-            value = self._rgb_left_value if not depth_flag else self._depth_left_value
+            queue = self._rgb_left_queue if not depth_flag else self._depth_left_queue
         elif camera_name == "right":
-            value = self._rgb_right_value if not depth_flag else self._depth_right_value
+            queue = self._rgb_right_queue if not depth_flag else self._depth_right_queue
         else:
             raise ValueError(f"unknown camera name: {camera_name}")
 
         try:
-            ts, count, img = value.get()
+            ts, count, img = 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

hex_zmq_servers/mujoco/mujoco_base.py

@@ -8,12 +8,12 @@
 
 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,
-    HexSafeValue,
     HexZMQClientBase,
     HexZMQServerBase,
 )
@@ -88,7 +88,7 @@ class HexMujocoBase(HexDeviceBase):
         return normed_rads
 
     @abstractmethod
-    def work_loop(self, hex_values: list[HexSafeValue | threading.Event]):
+    def work_loop(self, hex_queues: list[deque | threading.Event]):
         raise NotImplementedError(
             "`work_loop` should be implemented by the child class")
 
@@ -107,6 +107,13 @@ class HexMujocoClientBase(HexZMQClientBase):
         self._cmds_seq = 0
         self._rgb_seq = 0
         self._depth_seq = 0
+        self._states_queue = {
+            "robot": deque(maxlen=10),
+            "obj": deque(maxlen=10),
+        }
+        self._rgb_queue = deque(maxlen=10)
+        self._depth_queue = deque(maxlen=10)
+        self._cmds_queue = deque(maxlen=10)
 
     def __del__(self):
         HexZMQClientBase.__del__(self)
@@ -133,7 +140,77 @@ class HexMujocoClientBase(HexZMQClientBase):
         _, limits = self.request({"cmd": "get_limits"})
         return limits
 
-    def get_states(self, robot_name: str | None = None):
+    def get_states(self, robot_name: str | None = None, newest: bool = False):
+        try:
+            return self._states_queue[robot_name].popleft(
+            ) if not newest else self._states_queue[robot_name][-1]
+        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:
+            return self._rgb_queue.popleft(
+            ) if not newest else self._rgb_queue[-1]
+        except IndexError:
+            return None, None
+
+    def get_depth(self, camera_name: str | None = None, newest: bool = False):
+        try:
+            return self._depth_queue.popleft(
+            ) if not newest else self._depth_queue[-1]
+        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._depth_seq if depth_flag else self._rgb_seq)) %
+            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}"
@@ -156,7 +233,7 @@ class HexMujocoClientBase(HexZMQClientBase):
             print(f"\033[91m{req_cmd} failed: {e}\033[0m")
             return None, None
 
-    def set_cmds(
+    def _set_cmds_inner(
         self,
         cmds: np.ndarray,
     ) -> bool:
@@ -184,71 +261,28 @@ class HexMujocoClientBase(HexZMQClientBase):
             print(f"\033[91m{req_cmd} failed: {e}\033[0m")
             return False
 
-    def get_intri(self):
-        intri_hdr, intri = self.request({"cmd": "get_intri"})
-        return intri_hdr, intri
-
-    def get_rgb(self, camera_name: str | None = None):
-        return self._process_frame(camera_name, False)
-
-    def get_depth(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._depth_seq if depth_flag else self._rgb_seq)) %
-            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
-
 
 class HexMujocoServerBase(HexZMQServerBase):
 
     def __init__(self, net_config: dict = NET_CONFIG):
         HexZMQServerBase.__init__(self, net_config)
         self._device: HexDeviceBase = None
-        self._states_value = HexSafeValue()
-        self._obj_pose_value = HexSafeValue()
-        self._cmds_value = HexSafeValue()
+        self._states_queue = deque(maxlen=10)
+        self._obj_pose_queue = deque(maxlen=10)
+        self._cmds_queue = deque(maxlen=10)
         self._cmds_seq = -1
-        self._rgb_value = HexSafeValue()
-        self._depth_value = HexSafeValue()
+        self._rgb_queue = deque(maxlen=10)
+        self._depth_queue = deque(maxlen=10)
         self._seq_clear_flag = False
 
     def work_loop(self):
         try:
             self._device.work_loop([
-                self._states_value,
-                self._obj_pose_value,
-                self._cmds_value,
-                self._rgb_value,
-                self._depth_value,
+                self._states_queue,
+                self._obj_pose_queue,
+                self._cmds_queue,
+                self._rgb_queue,
+                self._depth_queue,
                 self._stop_event,
             ])
         finally:
@@ -266,7 +300,9 @@ class HexMujocoServerBase(HexZMQServerBase):
             return {"cmd": f"{recv_hdr['cmd']}_failed"}, None
 
         try:
-            ts, count, states = self._states_value.get()
+            ts, count, states = 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
@@ -292,7 +328,7 @@ class HexMujocoServerBase(HexZMQServerBase):
             delta = (seq - self._cmds_seq) % self._max_seq_num
             if delta >= 0 and delta < 1e6:
                 self._cmds_seq = seq
-                self._cmds_value.set((recv_hdr["ts"], seq, recv_buf))
+                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
@@ -309,13 +345,14 @@ class HexMujocoServerBase(HexZMQServerBase):
         # get camera config
         split_cmd = recv_hdr["cmd"].split("_")
         depth_flag = split_cmd[1] == "depth"
-        if depth_flag:
-            value = self._depth_value
-        else:
-            value = self._rgb_value
 
         try:
-            ts, count, img = value.get()
+            if depth_flag:
+                ts, count, img = self._depth_queue.popleft()
+            else:
+                ts, count, img = 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

hex_zmq_servers/robot/dummy/robot_dummy.py

@@ -8,13 +8,13 @@
 
 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,
-    HexSafeValue,
 )
 from ...hex_launch import hex_log, HEX_LOG_LEVEL
 
@@ -47,10 +47,10 @@ class HexRobotDummy(HexRobotBase):
     def __del__(self):
         HexRobotBase.__del__(self)
 
-    def work_loop(self, hex_values: list[HexSafeValue | threading.Event]):
-        states_value = hex_values[0]
-        cmds_value = hex_values[1]
-        stop_event = hex_values[2]
+    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
@@ -58,11 +58,15 @@ class HexRobotDummy(HexRobotBase):
         rate = HexRate(1000)
         while self._working.is_set() and not stop_event.is_set():
             # states
-            states_value.set((hex_zmq_ts_now(), states_count, dummy_states))
+            states_queue.append((hex_zmq_ts_now(), states_count, dummy_states))
             states_count = (states_count + 1) % self._max_seq_num
 
             # cmds
-            cmds_pack = cmds_value.get(timeout_s=-1.0)
+            cmds_pack = None
+            try:
+                cmds_pack = cmds_queue.popleft()
+            except IndexError:
+                pass
             if cmds_pack is not None:
                 ts, seq, cmds = cmds_pack
                 if seq != last_cmds_seq:

hex_zmq_servers/robot/dummy/robot_dummy_cli.py

@@ -24,3 +24,4 @@ class HexRobotDummyClient(HexRobotClientBase):
         net_config: dict = NET_CONFIG,
     ):
         HexRobotClientBase.__init__(self, net_config)
+        self._wait_for_working()