abizovnuralem / go2_ros2_sdk

Unofficial ROS2 SDK support for Unitree GO2 AIR/PRO/EDU
BSD 2-Clause "Simplified" License
231 stars 43 forks source link

Cannot control go2 when the second time. #52

Closed Lin-jun-xiang closed 2 months ago

Lin-jun-xiang commented 2 months ago

I extracted the WebRTC part for controlling Go2 (without ROS) but encountered the following issues:

webrtc_driver.py


import base64
import hashlib
import json
import logging
import time

import aiohttp
from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaBlackhole

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Go2Connection():
    def __init__(self, robot_ip, robot_num, token="", on_validated=None, on_message=None, on_open=None):
        self.pc = RTCPeerConnection()
        self.robot_ip = robot_ip
        self.robot_num = str(robot_num)
        self.token = token
        self.robot_validation = "PENDING"
        self.on_validated = on_validated
        self.on_message = on_message
        self.on_open = on_open

        self.audio_track = MediaBlackhole()
        self.video_track = MediaBlackhole()

        self.data_channel = self.pc.createDataChannel("data", id=0)
        self.data_channel.on("open", self.on_data_channel_open)
        self.data_channel.on("message", self.on_data_channel_message)

        self.pc.on("track", self.on_track)
        self.pc.on("connectionstatechange", self.on_connection_state_change)

    def on_connection_state_change(self):
        logger.info(f"Connection state is {self.pc.connectionState}")

    def on_track(self, track):
        logger.info(f"Receiving {track.kind}")

    async def generate_offer(self):
        await self.audio_track.start()
        await self.video_track.start()
        offer = await self.pc.createOffer()
        await self.pc.setLocalDescription(offer)
        return offer.sdp

    async def set_answer(self, sdp):
        answer = RTCSessionDescription(sdp, type="answer")
        await self.pc.setRemoteDescription(answer)

    def on_data_channel_open(self):
        logger.info("Data channel is open")
        if self.on_open:
            self.on_open()

    def on_data_channel_message(self, msg):
        logger.info(f"Received message: {msg}")

        if self.data_channel.readyState != "open":
            self.data_channel._setReadyState("open")

        try:
            if isinstance(msg, str):
                msgobj = json.loads(msg)
                if msgobj.get("type") == "validation":
                    self.validate_robot_conn(msgobj)
            if self.on_message:
                self.on_message(msg, msgobj, self.robot_num)
        except json.JSONDecodeError:
            pass

    async def connect(self):
        offer = await self.generate_offer()
        url = f"http://{self.robot_ip}:8081/offer"
        headers = {"Content-Type": "application/json"}
        data = {
            "sdp": offer,
            "id": "STA_localNetwork",
            "type": "offer",
            "token": "",
        }

        connected = False
        while not connected:
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=data, headers=headers) as resp:
                    if resp.status == 200:
                        answer_data = await resp.json()
                        answer_sdp = answer_data.get("sdp")
                        await self.set_answer(answer_sdp)
                        connected = True
                    else:
                        logger.info(f"Failed to get answer from server: Reason: {resp}")
                        logger.info("Try to reconnect...")
                        time.sleep(1)

    def validate_robot_conn(self, message):
        if message.get("data") == "Validation Ok.":
            self.robot_validation = "SUCCESS"
            if self.on_validated:
                self.on_validated(self.robot_num)
        else:
            self.publish(
                "",
                self.encrypt_key(message.get("data")),
                "validation",
            )

    def publish(self, topic, data, msg_type):
        if self.data_channel.readyState != "open":
            logger.info(f"Data channel is not open. State is {self.data_channel.readyState}")
            return

        payload = {
            "type": msg_type,
            "topic": topic,
            "data": data,
        }

        payload_dumped = json.dumps(payload)
        logger.info(f"-> Sending message {payload_dumped}")
        self.data_channel.send(payload_dumped)

    @staticmethod
    def hex_to_base64(hex_str):
        bytes_array = bytes.fromhex(hex_str)
        return base64.b64encode(bytes_array).decode("utf-8")

    @staticmethod
    def encrypt_key(key):
        prefixed_key = f"UnitreeGo2_{key}"
        encrypted = Go2Connection.encrypt_by_md5(prefixed_key)
        return Go2Connection.hex_to_base64(encrypted)

    @staticmethod
    def encrypt_by_md5(input_str):
        hash_obj = hashlib.md5()
        hash_obj.update(input_str.encode("utf-8"))
        return hash_obj.hexdigest()
import asyncio
import logging
import time

from scripts.go2_func import gen_command
from scripts.webrtc_driver import Go2Connection

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

robot_ip = "myip"
token = ""
robot_num = "0"

conn = Go2Connection(
    robot_ip=robot_ip,
    robot_num=robot_num,
    token=token,
    on_validated=lambda _: logger.info("Robot validated"),
    on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
)

async def connect():
    await conn.connect()
    logger.info("Connected to robot")

    while conn.robot_validation == "PENDING":
        await asyncio.sleep(0.1)

    if conn.robot_validation != "SUCCESS":
        logger.error("Failed to validate robot connection")
        return

    logger.info("Robot connection validated successfully")

def fun1():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun1 called')

def fun2():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun2 called')

if __name__ == "__main__":
    asyncio.run(connect())
    time.sleep(3)
    fun1()
    time.sleep(10)
    fun2()
import asyncio
import logging
import time

import nest_asyncio
from scripts.go2_func import gen_command
from scripts.webrtc_driver import Go2Connection

nest_asyncio.apply()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

robot_ip = "myip"
token = ""
robot_num = "0"

conn = Go2Connection(
    robot_ip=robot_ip,
    robot_num=robot_num,
    token=token,
    on_validated=lambda _: logger.info("Robot validated"),
    on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
)

async def connect():
    await conn.connect()
    logger.info("Connected to robot")

    while conn.robot_validation == "PENDING":
        await asyncio.sleep(0.1)

    if conn.robot_validation != "SUCCESS":
        logger.error("Failed to validate robot connection")
        return

    logger.info("Robot connection validated successfully")

def fun1():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun1 called')

def fun2():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun2 called')

if __name__ == "__main__":
    asyncio.run(connect())
    time.sleep(3)
    fun1()
    time.sleep(10)
    fun2()
fun1 called.
fun2 called.
import asyncio
import logging
import time
from functools import wraps

import nest_asyncio
from scripts.go2_func import gen_command
from scripts.webrtc_driver import Go2Connection

nest_asyncio.apply()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

robot_ip = "myip"
token = ""
robot_num = "0"

conn = Go2Connection(
    robot_ip=robot_ip,
    robot_num=robot_num,
    token=token,
    on_validated=lambda _: logger.info("Robot validated"),
    on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
)

async def connect():
    await conn.connect()
    logger.info("Connected to robot")

    while conn.robot_validation == "PENDING":
        await asyncio.sleep(0.1)

    if conn.robot_validation != "SUCCESS":
        logger.error("Failed to validate robot connection")
        return

    logger.info("Robot connection validated successfully")

def ensure_connected(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        if conn.pc.connectionState != "connected":
            logger.info("Robot not connected")
            return
        if conn.data_channel.readyState != "open":
            logger.info("Data channel not opened")
            return
        return await func(*args, **kwargs)
    return wrapper

@ensure_connected
async def fun1():
    cmd = gen_command(1029)
    logger.info(f'Sending command from fun1: {cmd}')
    conn.data_channel.send(cmd)
    logger.info('fun1 called')

@ensure_connected
async def fun2():
    cmd = gen_command(1029)
    logger.info(f'Sending command from fun2: {cmd}')
    conn.data_channel.send(cmd)
    logger.info('fun2 called')

if __name__ == "__main__":
    asyncio.run(connect())
    time.sleep(3)
    asyncio.run(fun1())
    time.sleep(10)
    asyncio.run(fun2())
fun1 called.
fun2 called.
import asyncio
import logging

from scripts.go2_constants import ROBOT_CMD
from scripts.go2_func import gen_command, gen_mov_command
from scripts.webrtc_driver import Go2Connection

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def main():
    robot_ip = "myip"
    token = ""
    robot_num = "0"

    conn = Go2Connection(
        robot_ip=robot_ip,
        robot_num=robot_num,
        token=token,
        on_validated=lambda _: logger.info("Robot validated"),
        on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
    )
    await conn.connect()
    logger.info("Connected to robot")

    while conn.robot_validation == "PENDING":
        await asyncio.sleep(0.1)

    if conn.robot_validation != "SUCCESS":
        logger.error("Failed to validate robot connection")
        return

    logger.info("Robot connection validated successfully")

    cmd = gen_command(1029)
    conn.data_channel.send(cmd)

    await asyncio.sleep(10)

    cmd = gen_command(1029)
    conn.data_channel.send(cmd)

    await asyncio.sleep(5)

    await conn.pc.close()
    logger.info("Disconnected from robot to stop movement")

if __name__ == "__main__":
    asyncio.run(main())
import asyncio
import logging
import nest_asyncio
from functools import wraps
from scripts.go2_func import gen_command
from scripts.webrtc_driver import Go2Connection

nest_asyncio.apply()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

robot_ip = "myip"
token = ""
robot_num = "0"

conn = Go2Connection(
    robot_ip=robot_ip,
    robot_num=robot_num,
    token=token,
    on_validated=lambda _: logger.info("Robot validated"),
    on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
)

async def connect():
    await conn.connect()
    logger.info("Connected to robot")

    while conn.robot_validation == "PENDING":
        await asyncio.sleep(0.1)

    if conn.robot_validation != "SUCCESS":
        logger.error("Failed to validate robot connection")
        return

    logger.info("Robot connection validated successfully")

def ensure_connected(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        if conn.pc.connectionState != "connected":
            logger.info("Robot not connected")
            return
        if conn.data_channel.readyState != "open":
            logger.info("Data channel not opened")
            return
        return await func(*args, **kwargs)
    return wrapper

@ensure_connected
async def fun1():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun1 called')

@ensure_connected
async def fun2():
    cmd = gen_command(1029)
    conn.data_channel.send(cmd)
    print('fun2 called')

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(connect())
    loop.create_task(fun1())
    loop.create_task(fun2())
    loop.run_forever()

What might be causing this issue? How can I ensure that different functions are called correctly without placing all commands in the same coroutine?

Lin-jun-xiang commented 2 months ago

@abizovnuralem Maybe the data send not correctly complete event loop in first time call fun?

abizovnuralem commented 2 months ago

Hi, there is no problem controlling the dog in real-time with the joystick from PC, something is wrong with the data that you are sending to it.

abizovnuralem commented 2 months ago

Also if you don't provide a conn token to fun1 and fun2 , you only will have 1 connection running at the same time.

Lin-jun-xiang commented 2 months ago

@abizovnuralem

Thanks your reply.

Is the conn token can be any like the following?

conn = Go2Connection(
    robot_ip=ROBOT_IP,
    robot_num="0",
    token="Go2",
    on_open=lambda: logger.info("Data channel opened"),
    on_validated=lambda _: logger.info("Robot validated"),
    on_message=lambda msg, _, __: logger.info(f"Message received: {msg}")
)

async def go2_scrape():
    conn.data_channel.send(
        gen_command(ROBOT_CMD['Scrape'])
    )
abizovnuralem commented 2 months ago

Obtaining security token Connecting to your device without a security token is possible and might allow a connection to be established. However, this method limits you to a single active connection at any time. To simultaneously use multiple clients, such as a WebRTC-based application and a phone app, a valid security token is necessary. This ensures secure, multi-client access to your device.

One way is to sniff the traffic between the dog and your phone. Assuming that you have Linux or Mac:

Run tinyproxy or any other HTTP proxy on your computer Set your computer's IP and port as HTTP proxy on your phone Run wireshark or ngrep on your box sniffing port 8081 like ngrep port 8081. Look for the token in the TCP stream after you connect your phone to the dog via the app The token looks like this in the request payload:

{ "token": "eyJ0eXAiOizI1NiJtlbiI[..]CI6MTcwODAxMzUwOX0.hiWOd9tNCIPzOOLNA", "sdp": "v=0\r\no=- ", "id": "STA_localNetwork", "type": "offer" } Another option is to obtain token via the /login/email endpoint.

curl -vX POST https://global-robot-api.unitree.com/login/email -d "email=&password="