controller_protocol.py

#!/usr/bin/env python3
"""
Unitree G1 Controller Protocol Implementation

Based on reverse engineering of unitree_sdk2:
https://github.com/unitreerobotics/unitree_sdk2

This module provides classes to:
1. Parse raw 40-byte controller packets
2. Decode button states and joystick values
3. Create virtual controller states for robot control
"""

import struct
from dataclasses import dataclass, field
from typing import Optional
from enum import IntFlag


class ControllerButtons(IntFlag):
    """Button bitmask mapping from Unitree controller"""
    R1 = 1 << 0
    L1 = 1 << 1
    START = 1 << 2
    SELECT = 1 << 3
    R2 = 1 << 4
    L2 = 1 << 5
    F1 = 1 << 6
    F2 = 1 << 7
    A = 1 << 8
    B = 1 << 9
    X = 1 << 10
    Y = 1 << 11
    UP = 1 << 12
    RIGHT = 1 << 13
    DOWN = 1 << 14
    LEFT = 1 << 15


@dataclass
class ControllerState:
    """
    Represents the current state of the Unitree wireless controller.
    
    Joystick values range from -1.0 to 1.0
    Trigger values range from 0.0 to 1.0
    """
    # Joysticks
    lx: float = 0.0  # Left stick X axis
    ly: float = 0.0  # Left stick Y axis
    rx: float = 0.0  # Right stick X axis
    ry: float = 0.0  # Right stick Y axis
    
    # Triggers
    l2: float = 0.0  # Left trigger
    r2: float = 0.0  # Right trigger (not in raw data, derived from button)
    
    # Buttons (raw bitmask)
    buttons: int = 0
    
    # Header bytes
    head: bytes = field(default_factory=lambda: b'\x00\x00')
    
    @property
    def r1(self) -> bool:
        return bool(self.buttons & ControllerButtons.R1)
    
    @property
    def l1(self) -> bool:
        return bool(self.buttons & ControllerButtons.L1)
    
    @property
    def start(self) -> bool:
        return bool(self.buttons & ControllerButtons.START)
    
    @property
    def select(self) -> bool:
        return bool(self.buttons & ControllerButtons.SELECT)
    
    @property
    def r2_button(self) -> bool:
        return bool(self.buttons & ControllerButtons.R2)
    
    @property
    def l2_button(self) -> bool:
        return bool(self.buttons & ControllerButtons.L2)
    
    @property
    def f1(self) -> bool:
        return bool(self.buttons & ControllerButtons.F1)
    
    @property
    def f2(self) -> bool:
        return bool(self.buttons & ControllerButtons.F2)
    
    @property
    def a(self) -> bool:
        return bool(self.buttons & ControllerButtons.A)
    
    @property
    def b(self) -> bool:
        return bool(self.buttons & ControllerButtons.B)
    
    @property
    def x(self) -> bool:
        return bool(self.buttons & ControllerButtons.X)
    
    @property
    def y(self) -> bool:
        return bool(self.buttons & ControllerButtons.Y)
    
    @property
    def up(self) -> bool:
        return bool(self.buttons & ControllerButtons.UP)
    
    @property
    def down(self) -> bool:
        return bool(self.buttons & ControllerButtons.DOWN)
    
    @property
    def left(self) -> bool:
        return bool(self.buttons & ControllerButtons.LEFT)
    
    @property
    def right(self) -> bool:
        return bool(self.buttons & ControllerButtons.RIGHT)
    
    def get_pressed_buttons(self) -> list[str]:
        """Return list of currently pressed button names"""
        pressed = []
        for btn in ControllerButtons:
            if self.buttons & btn:
                pressed.append(btn.name)
        return pressed
    
    def to_bytes(self) -> bytes:
        """
        Serialize controller state to 40-byte packet.
        
        Format:
        - 2 bytes: header
        - 2 bytes: button bitmask (uint16 LE)
        - 4 bytes: lx (float LE)
        - 4 bytes: rx (float LE)
        - 4 bytes: ry (float LE)
        - 4 bytes: l2 (float LE)
        - 4 bytes: ly (float LE)
        - 16 bytes: padding
        """
        data = struct.pack(
            '<2sHfffff16s',
            self.head,
            self.buttons & 0xFFFF,
            self.lx,
            self.rx,
            self.ry,
            self.l2,
            self.ly,
            b'\x00' * 16
        )
        return data
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'ControllerState':
        """
        Parse 40-byte packet into ControllerState.
        """
        if len(data) < 40:
            raise ValueError(f"Expected 40 bytes, got {len(data)}")
        
        head, buttons, lx, rx, ry, l2, ly, _ = struct.unpack(
            '<2sHfffff16s',
            data[:40]
        )
        
        return cls(
            head=head,
            buttons=buttons,
            lx=lx,
            ly=ly,
            rx=rx,
            ry=ry,
            l2=l2
        )
    
    def __str__(self) -> str:
        buttons_str = ', '.join(self.get_pressed_buttons()) or 'None'
        return (
            f"Controller State:\n"
            f"  Left Stick:  ({self.lx:+.3f}, {self.ly:+.3f})\n"
            f"  Right Stick: ({self.rx:+.3f}, {self.ry:+.3f})\n"
            f"  L2 Trigger:  {self.l2:.3f}\n"
            f"  Buttons:     {buttons_str}"
        )


class WirelessControllerMessage:
    """
    Represents the DDS WirelessController_ message format.
    
    This is the higher-level format used in the SDK's DDS communication,
    as opposed to the raw 40-byte serial format.
    
    Fields:
    - lx: float (left stick X)
    - ly: float (left stick Y)  
    - rx: float (right stick X)
    - ry: float (right stick Y)
    - keys: uint16 (button bitmask)
    """
    
    def __init__(self, lx: float = 0.0, ly: float = 0.0, 
                 rx: float = 0.0, ry: float = 0.0, keys: int = 0):
        self.lx = lx
        self.ly = ly
        self.rx = rx
        self.ry = ry
        self.keys = keys
    
    @classmethod
    def from_controller_state(cls, state: ControllerState) -> 'WirelessControllerMessage':
        """Convert from raw ControllerState to DDS message format"""
        return cls(
            lx=state.lx,
            ly=state.ly,
            rx=state.rx,
            ry=state.ry,
            keys=state.buttons
        )
    
    def to_controller_state(self) -> ControllerState:
        """Convert from DDS message format to ControllerState"""
        return ControllerState(
            lx=self.lx,
            ly=self.ly,
            rx=self.rx,
            ry=self.ry,
            buttons=self.keys
        )


# Movement command mapping based on SDK examples
class MovementCommands:
    """
    Standard movement mappings based on Unitree SDK examples.
    
    From user_controller.hpp:
    - cmd[0] = gamepad.ly   (forward/backward velocity)
    - cmd[1] = -gamepad.lx  (left/right velocity)
    - cmd[2] = -gamepad.rx  (yaw rotation)
    """
    
    @staticmethod
    def get_velocity_command(state: ControllerState) -> tuple[float, float, float]:
        """
        Convert controller state to velocity command.
        
        Returns:
            (forward_velocity, side_velocity, yaw_velocity)
        """
        return (
            state.ly,   # Forward/backward
            -state.lx,  # Left/right (strafe)
            -state.rx   # Yaw rotation
        )


def create_test_state() -> ControllerState:
    """Create a test controller state for development"""
    state = ControllerState()
    state.lx = 0.5
    state.ly = -0.3
    state.buttons = ControllerButtons.A | ControllerButtons.START
    return state


if __name__ == "__main__":
    # Demo usage
    print("=== Unitree Controller Protocol Demo ===\n")
    
    # Create a test state
    state = create_test_state()
    print(state)
    print()
    
    # Serialize to bytes
    packet = state.to_bytes()
    print(f"Serialized packet ({len(packet)} bytes):")
    print(f"  {packet.hex()}")
    print()
    
    # Parse back
    parsed = ControllerState.from_bytes(packet)
    print("Parsed back:")
    print(parsed)
    print()
    
    # Get velocity command
    vel = MovementCommands.get_velocity_command(state)
    print(f"Velocity command: forward={vel[0]:.2f}, side={vel[1]:.2f}, yaw={vel[2]:.2f}")