File size: 1,396 Bytes
301b0eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | """Exercise real wire packing, server handling, policy and mock IO without sockets/weights/hardware."""
from pathlib import Path
import asyncio,json,sys
sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'flexiv-umi-infer'))
from umi_infer.protocol import pack,unpack_action
from umi_infer.policies.mock import MockPolicy
from umi_infer.robot_io.mock import MockRobotIO
from umi_infer.policy_server import _handle
class InMemorySocket:
remote_address=('in-process',0)
def __init__(self,requests):self.requests=requests;self.responses=[]
def __aiter__(self):
async def messages():
for request in self.requests:yield request
return messages()
async def send(self,message):self.responses.append(message)
async def main():
io=MockRobotIO(image_hw=(32,32));io.start()
try:
for _ in range(3):
obs=io.get_observation();sock=InMemorySocket([pack(obs)])
await _handle(sock,MockPolicy(mode='hold'))
assert len(sock.responses)==1
action=unpack_action(sock.responses[0]);assert action.seq==obs.seq
assert set(action.eef_target)=={'left','right'}
io.apply_action(action)
print(json.dumps({'status':'passed','roundtrips':3,'transport':'in-process msgpack','robot':'mock','weights_loaded':False}))
finally:io.stop()
if __name__=='__main__':asyncio.run(main())
|