silencht
6 months ago
5 changed files with 644 additions and 168 deletions
-
4teleop/robot_control/robot_arm.py
-
5teleop/robot_control/robot_hand_unitree.py
-
375teleop/teleop_hand_and_arm.py
-
27teleop/utils/episode_writer.py
-
419teleop/utils/ipc.py
@ -0,0 +1,419 @@ |
|||
import os |
|||
import zmq |
|||
import time |
|||
import threading |
|||
import logging_mp |
|||
logging_mp.basic_config(level=logging_mp.INFO) |
|||
logger_mp = logging_mp.get_logger(__name__) |
|||
|
|||
""" |
|||
# Client → Server (Request) |
|||
1) launch |
|||
{ |
|||
"reqid": unique id, |
|||
"cmd": "CMD_START" |
|||
} |
|||
|
|||
2) exit |
|||
{ |
|||
"reqid": unique id, |
|||
"cmd": "CMD_STOP" |
|||
} |
|||
|
|||
3) start or stop (record toggle) |
|||
{ |
|||
"reqid": unique id, |
|||
"cmd": "CMD_RECORD_TOGGLE", |
|||
"info": { # optional |
|||
"task_name": "T001", |
|||
"task_desc": "pick and place apple to basket", |
|||
"item_id": 1 |
|||
} |
|||
} |
|||
|
|||
# Server → Client (Reply) |
|||
1) if ok |
|||
{ |
|||
"repid": same as reqid, |
|||
"status": "ok", |
|||
"msg": "ok" |
|||
} |
|||
2) if error |
|||
{ |
|||
"repid": same as reqid | 0 | 1, # 0: no reqid provided, 1: internal error |
|||
"status": "error", |
|||
"msg": "reqid not provided" |
|||
| "cmd not provided" |
|||
| "cmd not supported: {cmd}" |
|||
| "info missing keys: {missing_keys}" |
|||
| "internal error msg" |
|||
} |
|||
|
|||
# Heartbeat (PUB) |
|||
- Heartbeat Pub format: |
|||
{ |
|||
"START": True | False, # whether robot follow vr |
|||
"STOP" : True | False, # whether exit program |
|||
"RECORD_RUNNING": True | False, # whether is recording |
|||
"RECORD_READY": True | False, # whether ready to record |
|||
} |
|||
""" |
|||
|
|||
|
|||
class IPC_Server: |
|||
""" |
|||
Inter - Process Communication Server: |
|||
- Handle data via REP |
|||
- Publish heartbeat via PUB, Heartbeat state is provided by external callback get_state() |
|||
""" |
|||
# Mapping table for on_press keys |
|||
cmd_map = { |
|||
"CMD_START": "r", # launch |
|||
"CMD_STOP": "q", # exit |
|||
"CMD_RECORD_TOGGLE": "s", # start & stop (toggle record) |
|||
} |
|||
|
|||
def __init__(self, on_press=None, on_info=None, get_state=None, hb_fps=10.0): |
|||
""" |
|||
Args: |
|||
on_press : callback(cmd:str), called for every command |
|||
on_info : callback(data:dict), only handle CMD_RECORD_TOGGLE's task info |
|||
hb_fps : heartbeat publish frequency |
|||
get_state : callback() -> dict, provides current heartbeat state |
|||
""" |
|||
if callable(on_press): |
|||
self.on_press = on_press |
|||
else: |
|||
raise ValueError("[IPC_Server] on_press callback function must be provided") |
|||
self.on_info = on_info |
|||
if callable(get_state): |
|||
self.get_state = get_state |
|||
else: |
|||
raise ValueError("[IPC_Server] get_state callback function must be provided") |
|||
self._hb_interval = 1.0 / float(hb_fps) |
|||
self._running = True |
|||
self._data_loop_thread = None |
|||
self._hb_loop_thread = None |
|||
|
|||
rd = os.environ.get("XDG_RUNTIME_DIR") or "/tmp" |
|||
self.ctx = zmq.Context.instance() |
|||
# data IPC (REQ/REP): required |
|||
self.data_ipc = os.path.join(rd, f"xr-teleoperate-data-{os.getuid()}.ipc") |
|||
self.rep_socket = self.ctx.socket(zmq.REP) |
|||
try: |
|||
if os.path.exists(self.data_ipc): |
|||
os.unlink(self.data_ipc) # remove stale IPC file |
|||
except OSError: |
|||
pass |
|||
self.rep_socket.bind(f"ipc://{self.data_ipc}") |
|||
logger_mp.info(f"[IPC_Server] Listening to Data at ipc://{self.data_ipc}") |
|||
|
|||
# heartbeat IPC (PUB/SUB) |
|||
self.hb_ipc = os.path.join(rd, f"xr-teleoperate-hb-{os.getuid()}.ipc") |
|||
self.pub_socket = self.ctx.socket(zmq.PUB) |
|||
try: |
|||
if os.path.exists(self.hb_ipc): |
|||
os.unlink(self.hb_ipc) # remove stale IPC file |
|||
except OSError: |
|||
pass |
|||
self.pub_socket.bind(f"ipc://{self.hb_ipc}") |
|||
logger_mp.info(f"[IPC_Server] Publishing HeartBeat at ipc://{self.hb_ipc}") |
|||
|
|||
def _data_loop(self): |
|||
""" |
|||
Listen for REQ/REP commands and optional info. |
|||
""" |
|||
poller = zmq.Poller() |
|||
poller.register(self.rep_socket, zmq.POLLIN) |
|||
while self._running: |
|||
try: |
|||
socks = dict(poller.poll(20)) |
|||
if self.rep_socket in socks: |
|||
msg = self.rep_socket.recv_json() |
|||
reply = self._handle_message(msg) |
|||
try: |
|||
self.rep_socket.send_json(reply) |
|||
except Exception as e: |
|||
logger_mp.error(f"[IPC_Server] Failed to send reply: {e}") |
|||
finally: |
|||
logger_mp.debug(f"[IPC_Server] DATA recv: {msg} -> rep: {reply}") |
|||
except zmq.error.ContextTerminated: |
|||
break |
|||
except Exception as e: |
|||
logger_mp.error(f"[IPC_Server] Data loop exception: {e}") |
|||
|
|||
def _hb_loop(self): |
|||
"""Publish heartbeat periodically""" |
|||
while self._running: |
|||
start_time = time.monotonic() |
|||
try: |
|||
state = dict(self.get_state() or {}) |
|||
self.pub_socket.send_json(state) |
|||
logger_mp.debug(f"[IPC_Server] HB pub: {state}") |
|||
except Exception as e: |
|||
logger_mp.error(f"[IPC_Server] HeartBeat loop exception: {e}") |
|||
elapsed = time.monotonic() - start_time |
|||
if elapsed < self._hb_interval: |
|||
time.sleep(self._hb_interval - elapsed) |
|||
|
|||
def _handle_message(self, msg: dict) -> dict: |
|||
"""Process message and return reply""" |
|||
try: |
|||
# validate reqid |
|||
reqid = msg.get("reqid", None) |
|||
if not reqid: |
|||
return {"repid": 0, "status": "error", "msg": "reqid not provided"} |
|||
|
|||
# validate cmd |
|||
cmd = msg.get("cmd", None) |
|||
if not cmd: |
|||
return {"repid": reqid, "status": "error", "msg": "cmd not provided"} |
|||
|
|||
# unsupported cmd |
|||
if cmd not in self.cmd_map: |
|||
return {"repid": reqid, "status": "error", "msg": f"cmd not supported: {cmd}"} |
|||
|
|||
# CMD_RECORD_TOGGLE: optional info |
|||
if cmd == "CMD_RECORD_TOGGLE": |
|||
info = msg.get("info", None) |
|||
if info: |
|||
required_keys = ["task_name", "task_desc", "item_id"] |
|||
missing_keys = [key for key in required_keys if key not in info] |
|||
if missing_keys: |
|||
return {"repid": reqid, "status": "error", "msg": f"info missing keys: {missing_keys}"} |
|||
else: |
|||
if self.on_info: |
|||
self.on_info(info) |
|||
logger_mp.debug(f"[IPC_Server] on_info called with info: {info}") |
|||
else: |
|||
logger_mp.warning("[IPC_Server] No on_info provided") |
|||
else: |
|||
logger_mp.warning("[IPC_Server] No info provided with cmd: CMD_RECORD_TOGGLE") |
|||
|
|||
# supported cmd path |
|||
self.on_press(self.cmd_map[cmd]) |
|||
return {"repid": reqid, "status": "ok", "msg": "ok"} |
|||
|
|||
except Exception as e: |
|||
return {"repid": 1, "status": "error", "msg": str(e)} |
|||
|
|||
# --------------------------- |
|||
# Public API |
|||
# --------------------------- |
|||
def start(self): |
|||
"""Start both data loop and heartbeat loop""" |
|||
self._data_loop_thread = threading.Thread(target=self._data_loop, daemon=True) |
|||
self._data_loop_thread.start() |
|||
self._hb_loop_thread = threading.Thread(target=self._hb_loop, daemon=True) |
|||
self._hb_loop_thread.start() |
|||
|
|||
def stop(self): |
|||
"""Stop server""" |
|||
self._running = False |
|||
if self._data_loop_thread: |
|||
self._data_loop_thread.join(timeout=1.0) |
|||
if self._hb_loop_thread: |
|||
self._hb_loop_thread.join(timeout=1.0) |
|||
try: |
|||
self.rep_socket.setsockopt(zmq.LINGER, 0) |
|||
self.rep_socket.close() |
|||
except Exception: |
|||
pass |
|||
try: |
|||
self.pub_socket.setsockopt(zmq.LINGER, 0) |
|||
self.pub_socket.close() |
|||
except Exception: |
|||
pass |
|||
try: |
|||
self.ctx.term() |
|||
except Exception: |
|||
pass |
|||
|
|||
|
|||
class IPC_Client: |
|||
""" |
|||
Inter - Process Communication Client: |
|||
- Send command/info via REQ |
|||
- Subscribe heartbeat via SUB |
|||
""" |
|||
def __init__(self, hb_fps=10.0): |
|||
"""hb_fps: heartbeat subscribe frequency, should match server side.""" |
|||
rd = os.environ.get("XDG_RUNTIME_DIR") or "/tmp" |
|||
self.ctx = zmq.Context.instance() |
|||
|
|||
# heartbeat IPC (PUB/SUB) |
|||
self._hb_running = True |
|||
self._hb_last_time = 0 # timestamp of last heartbeat received |
|||
self._hb_latest_state = {} # latest heartbeat state |
|||
self._hb_online = False # whether heartbeat is online |
|||
self._hb_interval = 1.0 / float(hb_fps) # expected heartbeat interval |
|||
self._hb_lock = threading.Lock() # lock for heartbeat state |
|||
self._hb_timeout = 5.0 * self._hb_interval # timeout to consider offline |
|||
self.hb_ipc = os.path.join(rd, f"xr-teleoperate-hb-{os.getuid()}.ipc") |
|||
self.sub_socket = self.ctx.socket(zmq.SUB) |
|||
self.sub_socket.setsockopt(zmq.RCVHWM, 1) |
|||
self.sub_socket.connect(f"ipc://{self.hb_ipc}") |
|||
self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "") |
|||
logger_mp.info(f"[IPC_Client] Subscribed to HeartBeat at ipc://{self.hb_ipc}") |
|||
self._hb_thread = threading.Thread(target=self._hb_loop, daemon=True) |
|||
self._hb_thread.start() |
|||
|
|||
# data IPC (REQ/REP) |
|||
self.data_ipc = os.path.join(rd, f"xr-teleoperate-data-{os.getuid()}.ipc") |
|||
self.req_socket = self.ctx.socket(zmq.REQ) |
|||
self.req_socket.connect(f"ipc://{self.data_ipc}") |
|||
logger_mp.info(f"[IPC_Client] Connected to Data at ipc://{self.data_ipc}") |
|||
|
|||
def _make_reqid(self) -> str: |
|||
import uuid |
|||
return str(uuid.uuid4()) |
|||
|
|||
# --------------------------- |
|||
# Heartbeat handling |
|||
# --------------------------- |
|||
def _hb_loop(self): |
|||
consecutive = 0 |
|||
while self._hb_running: |
|||
start_time = time.monotonic() |
|||
try: |
|||
msg = self.sub_socket.recv_json(flags=zmq.NOBLOCK) |
|||
with self._hb_lock: |
|||
self._hb_latest_state = msg |
|||
self._hb_last_time = time.monotonic() |
|||
consecutive += 1 |
|||
if consecutive >= 3: # require 3 consecutive heartbeats to be considered online |
|||
self._hb_online = True |
|||
except zmq.Again: |
|||
with self._hb_lock: |
|||
if self._hb_last_time > 0: |
|||
if self._hb_online and (time.monotonic() - self._hb_last_time > self._hb_timeout): |
|||
self._hb_latest_state = {} |
|||
self._hb_last_time = 0 |
|||
self._hb_online = False |
|||
consecutive = 0 |
|||
logger_mp.warning("[IPC_Client] HeartBeat timeout -> OFFLINE") |
|||
except Exception as e: |
|||
logger_mp.error(f"[IPC_Client] HB loop exception: {e}") |
|||
elapsed = time.monotonic() - start_time |
|||
if elapsed < self._hb_interval: |
|||
time.sleep(self._hb_interval - elapsed) |
|||
|
|||
# --------------------------- |
|||
# Public API |
|||
# --------------------------- |
|||
def send_data(self, cmd: str, info: dict = None) -> dict: |
|||
"""Send command to server and wait reply""" |
|||
reqid = self._make_reqid() |
|||
if not self.is_online(): |
|||
logger_mp.warning(f"[IPC_Client] Cannot send {cmd}, server offline (no heartbeat)") |
|||
return {"repid": reqid, "status": "error", "msg": "server offline (no heartbeat)"} |
|||
|
|||
msg = {"reqid": reqid, "cmd": cmd} |
|||
if cmd == "CMD_RECORD_TOGGLE" and info: |
|||
msg["info"] = info |
|||
try: |
|||
self.req_socket.send_json(msg) |
|||
# wait up to 1s for reply |
|||
if self.req_socket.poll(1000): |
|||
reply = self.req_socket.recv_json() |
|||
else: |
|||
return {"repid": reqid, "status": "error", "msg": "timeout waiting for server reply"} |
|||
except Exception as e: |
|||
logger_mp.error(f"[IPC_Client] send_data failed: {e}") |
|||
return {"repid": reqid, "status": "error", "msg": str(e)} |
|||
|
|||
if reply.get("status") != "ok": |
|||
return reply |
|||
if reply.get("repid") != reqid: |
|||
return {"repid": reqid, "status": "error", "msg": f"reply id mismatch: expected {reqid}, got {reply.get('repid')}"} |
|||
return reply |
|||
|
|||
def is_online(self) -> bool: |
|||
with self._hb_lock: |
|||
return self._hb_online |
|||
|
|||
def latest_state(self) -> dict: |
|||
with self._hb_lock: |
|||
return dict(self._hb_latest_state) |
|||
|
|||
def stop(self): |
|||
self._hb_running = False |
|||
if self._hb_thread: |
|||
self._hb_thread.join(timeout=1.0) |
|||
try: |
|||
self.req_socket.setsockopt(zmq.LINGER, 0) |
|||
self.req_socket.close() |
|||
except Exception: |
|||
pass |
|||
try: |
|||
self.sub_socket.setsockopt(zmq.LINGER, 0) |
|||
self.sub_socket.close() |
|||
except Exception: |
|||
pass |
|||
try: |
|||
self.ctx.term() |
|||
except Exception: |
|||
pass |
|||
|
|||
|
|||
# --------------------------- |
|||
# Client Example usage |
|||
# --------------------------- |
|||
if __name__ == "__main__": |
|||
from sshkeyboard import listen_keyboard, stop_listening |
|||
client = None |
|||
|
|||
def on_press(key: str): |
|||
global client |
|||
if client is None: |
|||
logger_mp.warning("⚠️ Client not initialized, ignoring key press") |
|||
return |
|||
|
|||
if key == "r": |
|||
logger_mp.info("▶️ Sending launch command...") |
|||
rep = client.send_data("CMD_START") |
|||
logger_mp.info("Reply: %s", rep) |
|||
|
|||
elif key == "s": |
|||
info = { |
|||
"task_name": "T003", |
|||
"task_desc": "pick and place pear.", |
|||
"item_id": 1, |
|||
} |
|||
logger_mp.info("⏺️ Sending record toggle command...") |
|||
rep = client.send_data("CMD_RECORD_TOGGLE", info=info) # optional info |
|||
logger_mp.info("Reply: %s", rep) |
|||
|
|||
|
|||
elif key == "q": |
|||
logger_mp.info("⏹️ Sending exit command...") |
|||
rep = client.send_data("CMD_STOP") |
|||
logger_mp.info("Reply: %s", rep) |
|||
|
|||
elif key == "b": |
|||
if client.is_online(): |
|||
state = client.latest_state() |
|||
logger_mp.info(f"[HEARTBEAT] Current heartbeat: {state}") |
|||
else: |
|||
logger_mp.warning("[HEARTBEAT] No heartbeat received (OFFLINE)") |
|||
|
|||
else: |
|||
logger_mp.warning(f"⚠️ Undefined key: {key}") |
|||
|
|||
# Initialize client |
|||
client = IPC_Client(hb_fps=10.0) |
|||
|
|||
# Start keyboard listening thread |
|||
listen_keyboard_thread = threading.Thread(target=listen_keyboard, kwargs={"on_press": on_press, "until": None, "sequential": False}, daemon=True) |
|||
listen_keyboard_thread.start() |
|||
|
|||
logger_mp.info("✅ Client started, waiting for keyboard input:\n [r] launch, [s] start/stop record, [b] heartbeat, [q] exit") |
|||
try: |
|||
while True: |
|||
time.sleep(1.0) |
|||
except KeyboardInterrupt: |
|||
logger_mp.info("⏹️ User interrupt, preparing to exit...") |
|||
finally: |
|||
stop_listening() |
|||
client.stop() |
|||
logger_mp.info("✅ Client exited") |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue