hex-zmq-servers 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

hex_zmq_servers/mujoco/archer_y6/mujoco_archer_y6_cli.py

@@ -7,6 +7,7 @@
 ################################################################
 
 from ..mujoco_base import HexMujocoClientBase
+from ...zmq_base import HexRate
 
 NET_CONFIG = {
     "ip": "127.0.0.1",
@@ -16,11 +17,52 @@ NET_CONFIG = {
     "server_num_workers": 4,
 }
 
+RECV_CONFIG = {
+    "rgb": True,
+    "depth": True,
+    "obj": True,
+}
+
 
 class HexMujocoArcherY6Client(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._wait_for_working()
+
+    def _recv_loop(self):
+        rate = HexRate(2000)
+        image_trig_cnt = 0
+        while self._recv_flag:
+            hdr, states = self._get_states_inner("robot")
+            if hdr is not None:
+                self._states_queue["robot"].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["rgb"]:
+                    hdr, img = self._get_rgb_inner()
+                    if hdr is not None:
+                        self._rgb_queue.append((hdr, img))
+                if self.__recv_config["depth"]:
+                    hdr, img = self._get_depth_inner()
+                    if hdr is not None:
+                        self._depth_queue.append((hdr, img))
+
+            try:
+                cmds = self._cmds_queue.popleft()
+                _ = self._set_cmds_inner(cmds)
+            except IndexError:
+                pass
+
+            rate.sleep()

hex_zmq_servers/mujoco/archer_y6/mujoco_archer_y6_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,
@@ -51,19 +51,19 @@ class HexMujocoArcherY6Server(HexMujocoServerBase):
         self._device = HexMujocoArcherY6(params_config)
 
         # values
-        self._states_robot_value = HexSafeValue()
-        self._states_obj_value = HexSafeValue()
-        self._rgb_value = HexSafeValue()
-        self._depth_value = HexSafeValue()
+        self._states_robot_queue = deque(maxlen=10)
+        self._states_obj_queue = deque(maxlen=10)
+        self._rgb_queue = deque(maxlen=10)
+        self._depth_queue = deque(maxlen=10)
 
     def work_loop(self):
         try:
             self._device.work_loop([
-                self._states_robot_value,
-                self._states_obj_value,
-                self._cmds_value,
-                self._rgb_value,
-                self._depth_value,
+                self._states_robot_queue,
+                self._states_obj_queue,
+                self._cmds_queue,
+                self._rgb_queue,
+                self._depth_queue,
                 self._stop_event,
             ])
         finally:
@@ -79,15 +79,17 @@ class HexMujocoArcherY6Server(HexMujocoServerBase):
         # get robot name
         robot_name = recv_hdr["cmd"].split("_")[2]
         if robot_name == "robot":
-            value = self._states_robot_value
+            queue = self._states_robot_queue
         elif robot_name == "obj":
-            value = self._states_obj_value
+            queue = self._states_obj_queue
         else:
             raise ValueError(
                 f"unknown robot name: {robot_name} in {recv_hdr['cmd']}")
 
         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

hex_zmq_servers/mujoco/e3_desktop/model/scene.xml

@@ -1,7 +1,7 @@
 <mujoco model="e3_desktop scene">
     <compiler angle="radian"/>
     <option integrator="implicitfast"/>
-    <option timestep="2e-3"/>
+    <option timestep="1e-3"/>
 
     <statistic center="0.3 0 0.45" extent="0.8" meansize="0.05"/>
 

hex_zmq_servers/mujoco/e3_desktop/mujoco_e3_desktop.py

@@ -11,22 +11,24 @@ import copy
 import threading
 import cv2
 import numpy as np
+from collections import deque
 
 import mujoco
 from mujoco import viewer
 
 from ..mujoco_base import HexMujocoBase
 from ...zmq_base import (
+    hex_ns_now,
     hex_zmq_ts_now,
+    ns_to_hex_zmq_ts,
     hex_zmq_ts_delta_ms,
     HexRate,
-    HexSafeValue,
 )
 from ...hex_launch import hex_log, HEX_LOG_LEVEL
 from hex_robo_utils import HexCtrlUtilMitJoint as CtrlUtil
 
 MUJOCO_CONFIG = {
-    "states_rate": 500,
+    "states_rate": 1000,
     "img_rate": 30,
     "tau_ctrl": False,
     "mit_kp": [200.0, 200.0, 200.0, 75.0, 15.0, 15.0, 20.0],
@@ -146,6 +148,9 @@ class HexMujocoE3Desktop(HexMujocoBase):
         if not self.__headless:
             self.__viewer = viewer.launch_passive(self.__model, self.__data)
 
+        # time init
+        self.__bias_ns = hex_ns_now() - self.__data.time * 1_000_000_000
+
         # start work loop
         self._working.set()
 
@@ -161,19 +166,19 @@ class HexMujocoE3Desktop(HexMujocoBase):
             self.__viewer.sync()
         return True
 
-    def work_loop(self, hex_values: list[HexSafeValue | threading.Event]):
-        states_left_value = hex_values[0]
-        states_right_value = hex_values[1]
-        states_obj_value = hex_values[2]
-        cmds_left_value = hex_values[3]
-        cmds_right_value = hex_values[4]
-        head_rgb_value = hex_values[5]
-        head_depth_value = hex_values[6]
-        left_rgb_value = hex_values[7]
-        left_depth_value = hex_values[8]
-        right_rgb_value = hex_values[9]
-        right_depth_value = hex_values[10]
-        stop_event = hex_values[11]
+    def work_loop(self, hex_queues: list[deque | threading.Event]):
+        states_left_queue = hex_queues[0]
+        states_right_queue = hex_queues[1]
+        states_obj_queue = hex_queues[2]
+        cmds_left_queue = hex_queues[3]
+        cmds_right_queue = hex_queues[4]
+        head_rgb_queue = hex_queues[5]
+        head_depth_queue = hex_queues[6]
+        left_rgb_queue = hex_queues[7]
+        left_depth_queue = hex_queues[8]
+        right_rgb_queue = hex_queues[9]
+        right_depth_queue = hex_queues[10]
+        stop_event = hex_queues[11]
 
         last_states_ts = {"s": 0, "ns": 0}
         states_left_count = 0
@@ -194,22 +199,22 @@ class HexMujocoE3Desktop(HexMujocoBase):
         states_trig_count = 0
         img_trig_count = 0
         init_ts = self.__mujoco_ts() if self.__sens_ts else hex_zmq_ts_now()
-        head_rgb_value.set((init_ts, 0,
+        head_rgb_queue.append((init_ts, 0,
                             np.zeros((self.__height, self.__width, 3),
                                      dtype=np.uint8)))
-        head_depth_value.set((init_ts, 0,
+        head_depth_queue.append((init_ts, 0,
                               np.zeros((self.__height, self.__width),
                                        dtype=np.uint16)))
-        left_rgb_value.set((init_ts, 0,
+        left_rgb_queue.append((init_ts, 0,
                             np.zeros((self.__height, self.__width, 3),
                                      dtype=np.uint8)))
-        left_depth_value.set((init_ts, 0,
+        left_depth_queue.append((init_ts, 0,
                               np.zeros((self.__height, self.__width),
                                        dtype=np.uint16)))
-        right_rgb_value.set((init_ts, 0,
+        right_rgb_queue.append((init_ts, 0,
                              np.zeros((self.__height, self.__width, 3),
                                       dtype=np.uint8)))
-        right_depth_value.set((init_ts, 0,
+        right_depth_queue.append((init_ts, 0,
                                np.zeros((self.__height, self.__width),
                                         dtype=np.uint16)))
         while self._working.is_set() and not stop_event.is_set():
@@ -220,36 +225,44 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 # states
                 ts, states_left, states_right, states_obj = self.__get_states()
                 if states_left is not None:
-                    if hex_zmq_ts_delta_ms(ts, last_states_ts) > 1.0:
+                    if hex_zmq_ts_delta_ms(ts, last_states_ts) > 1e-6:
                         last_states_ts = ts
                         # states left
-                        states_left_value.set(
+                        states_left_queue.append(
                             (ts, states_left_count, states_left))
                         states_left_count = (states_left_count +
                                              1) % self._max_seq_num
                         # states right
-                        states_right_value.set(
+                        states_right_queue.append(
                             (ts, states_right_count, states_right))
                         states_right_count = (states_right_count +
                                               1) % self._max_seq_num
                         # states obj
-                        states_obj_value.set(
+                        states_obj_queue.append(
                             (ts, states_obj_count, states_obj))
                         states_obj_count = (states_obj_count +
                                             1) % self._max_seq_num
 
                 # cmds
-                cmds_left_pack = cmds_left_value.get(timeout_s=-1.0)
+                cmds_left_pack = None
+                try:
+                    cmds_left_pack = cmds_left_queue.popleft()
+                except IndexError:
+                    pass
                 if cmds_left_pack is not None:
-                    ts, seq, cmds_left_get = cmds_left_pack
+                    ts, seq, cmds_left = cmds_left_pack
                     if seq != last_cmds_left_seq:
                         last_cmds_left_seq = seq
                         if hex_zmq_ts_delta_ms(hex_zmq_ts_now(), ts) < 200.0:
-                            cmds_left = cmds_left_get.copy()
+                            cmds_left = cmds_left.copy()
                 if cmds_left is not None:
                     self.__set_cmds(cmds_left, "left")
 
-                cmds_right_pack = cmds_right_value.get(timeout_s=-1.0)
+                cmds_right_pack = None
+                try:
+                    cmds_right_pack = cmds_right_queue.popleft()
+                except IndexError:
+                    pass
                 if cmds_right_pack is not None:
                     ts, seq, cmds_right_get = cmds_right_pack
                     if seq != last_cmds_right_seq:
@@ -267,7 +280,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__head_rgb:
                     ts, rgb_img = self.__get_rgb("head_camera")
                     if rgb_img is not None:
-                        head_rgb_value.set((ts, head_rgb_count, rgb_img))
+                        head_rgb_queue.append((ts, head_rgb_count, rgb_img))
                         head_rgb_count = (head_rgb_count +
                                           1) % self._max_seq_num
 
@@ -275,7 +288,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__head_depth:
                     ts, depth_img = self.__get_depth("head_camera")
                     if depth_img is not None:
-                        head_depth_value.set((ts, head_depth_count, depth_img))
+                        head_depth_queue.append((ts, head_depth_count, depth_img))
                         head_depth_count = (head_depth_count +
                                             1) % self._max_seq_num
 
@@ -283,7 +296,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__left_rgb:
                     ts, rgb_img = self.__get_rgb("left_camera")
                     if rgb_img is not None:
-                        left_rgb_value.set((ts, left_rgb_count, rgb_img))
+                        left_rgb_queue.append((ts, left_rgb_count, rgb_img))
                         left_rgb_count = (left_rgb_count +
                                           1) % self._max_seq_num
 
@@ -291,7 +304,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__left_depth:
                     ts, depth_img = self.__get_depth("left_camera")
                     if depth_img is not None:
-                        left_depth_value.set((ts, left_depth_count, depth_img))
+                        left_depth_queue.append((ts, left_depth_count, depth_img))
                         left_depth_count = (left_depth_count +
                                             1) % self._max_seq_num
 
@@ -299,7 +312,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__right_rgb:
                     ts, rgb_img = self.__get_rgb("right_camera")
                     if rgb_img is not None:
-                        right_rgb_value.set((ts, right_rgb_count, rgb_img))
+                        right_rgb_queue.append((ts, right_rgb_count, rgb_img))
                         right_rgb_count = (right_rgb_count +
                                            1) % self._max_seq_num
 
@@ -307,7 +320,7 @@ class HexMujocoE3Desktop(HexMujocoBase):
                 if self.__right_depth:
                     ts, depth_img = self.__get_depth("right_camera")
                     if depth_img is not None:
-                        right_depth_value.set(
+                        right_depth_queue.append(
                             (ts, right_depth_count, depth_img))
                         right_depth_count = (right_depth_count +
                                              1) % self._max_seq_num
@@ -415,11 +428,8 @@ class HexMujocoE3Desktop(HexMujocoBase):
         ), depth_img
 
     def __mujoco_ts(self):
-        mujoco_ts = self.__data.time
-        return {
-            "s": int(mujoco_ts // 1),
-            "ns": int((mujoco_ts % 1) * 1_000_000_000),
-        }
+        mujoco_ts = self.__data.time * 1_000_000_000 + self.__bias_ns
+        return ns_to_hex_zmq_ts(mujoco_ts)
 
     def close(self):
         if not self._working.is_set():

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()