naowalk/naowalk.py

#!/usr/bin/env python
# -*- encoding: UTF-8 -*-

import qi
import argparse
import sys
import time
import pygame
import cv2
import numpy as np

class NaoTeleop:
    def __init__(self, session):
        self.session = session
        self.motion = session.service("ALMotion")
        self.posture = session.service("ALRobotPosture")
        
        # New: Hook into the battery service
        try:
            self.battery = session.service("ALBattery")
        except:
            self.battery = None
            print("⚠️ ALBattery service not found")

        self.video = None
        self.video_client = None
        self.running = True

        self.head_yaw = 0.0
        self.head_pitch = 0.0
        
        # Battery tracking
        self.battery_level = 100
        self.last_batt_check = 0

        self._init_video_maxfps()

        pygame.init()
        self.screen = pygame.display.set_mode((320, 240))
        pygame.display.set_caption("NAO Teleop | Battery: Checking...")
        self.clock = pygame.time.Clock()

    def _init_video_maxfps(self):
        try:
            self.video = self.session.service("ALVideoDevice")
            self.video_client = self.video.subscribeCamera("TeleopCam", 0, 0, 11, 25)
            print("✅ Camera ready")
        except:
            print("❌ Camera not available")

    def get_image(self):
        if not self.video or not self.video_client:
            return None
        try:
            image = self.video.getImageRemote(self.video_client)
            if image and len(image) >= 7:
                w, h = image[0], image[1]
                arr = np.frombuffer(bytearray(image[6]), dtype=np.uint8).reshape((h, w, 3))
                frame = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
                return cv2.resize(frame, (320, 240), interpolation=cv2.INTER_NEAREST)
        except:
            pass
        return None

    def wave_gesture(self):
        print("🤚 Waving hello...")
        try:
            # Stiffen arm for accurate wave
            self.motion.setStiffnesses("RArm", 1.0)
            
            # Use all 6 joints to prevent xAssertArraySize crash
            names = ["RShoulderPitch", "RShoulderRoll", "RElbowYaw", "RElbowRoll", "RWristYaw", "RHand"]
            
            angles = [
                [-1.5, -1.5, -1.5, -1.5, -1.5],  # RShoulderPitch (high up)
                [-0.4, -0.7, -0.4, -0.7, -0.4],  # RShoulderRoll
                [ 0.6,  1.1,  0.6,  1.1,  0.6],  # RElbowYaw
                [ 0.8,  0.3,  0.8,  0.3,  0.8],  # RElbowRoll (proper bending)
                [ 0.0,  1.5,  0.0, -1.5,  0.0],  # RWristYaw (the wave)
                [ 1.0,  1.0,  1.0,  1.0,  1.0]   # RHand (keep open)
            ]
            
            times = [[0.5, 1.0, 1.5, 2.0, 2.5]] * 6
            
            self.motion.angleInterpolation(names, angles, times, True)
            
            # Safely lower arm to true neutral
            neutral_angles = [1.5, -0.15, 1.2, 0.5, 0.0, 0.0]
            self.motion.angleInterpolationWithSpeed(names, neutral_angles, 0.3)
            
            # Relax the arm's stiffness to let the motors cool and walk naturally
            self.motion.setStiffnesses("RArm", 0.6)
            
            print("✅ Proper wave completed and arm relaxed")
            
        except Exception as e:
            print(f"Wave error: {e}")

    def handle_head_movement(self):
        keys = pygame.key.get_pressed()
        speed = 0.3
        changed = False

        if keys[pygame.K_i]:   
            self.head_pitch = max(self.head_pitch - speed, -0.5); changed = True
        if keys[pygame.K_k]:   
            self.head_pitch = min(self.head_pitch + speed, 0.5); changed = True
        if keys[pygame.K_j]:   
            self.head_yaw = min(self.head_yaw + speed, 2.0); changed = True
        if keys[pygame.K_l]:   
            self.head_yaw = max(self.head_yaw - speed, -2.0); changed = True

        if changed:
            self.motion.setAngles(["HeadYaw", "HeadPitch"], [self.head_yaw, self.head_pitch], 0.3)

    def check_battery(self):
        # Only ping the robot for battery data once every 10 seconds to save FPS
        if not self.battery:
            return
            
        now = time.time()
        if now - self.last_batt_check > 10:
            try:
                self.battery_level = self.battery.getBatteryCharge()
                pygame.display.set_caption(f"NAO Teleop | Battery: {self.battery_level}%")
            except:
                pass
            self.last_batt_check = now

    def run(self):
        self.motion.wakeUp()
        self.posture.goToPosture("StandInit", 0.5)
        
        # 1. SETTLE DELAY: Let his internal gyroscopes adapt to the squishy carpet
        time.sleep(1.0)
        
        self.motion.setMoveArmsEnabled(True, True)
        
        # 2. CARPET GAIT CONFIGURATION
        # StepHeight: Lifts the foot 2.8cm (higher than default) to clear fibers
        # MaxStepFrequency: 0.6 slows the step rhythm down for better balance
        # TorsoWy: 0.05 leans his torso just slightly forward for better momentum
        carpet_config = [
            ["StepHeight", 0.028],      
            ["MaxStepFrequency", 0.6],  
            ["TorsoWy", 0.05]           
        ]
        
        print("🎮 Controls:")
        print("   WASD / Arrows = Walk")
        print("   I J K L       = Head")
        print("   1             = Wave")
        print("   ESC = Quit")

        while self.running:
            self.check_battery()

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        self.running = False
                    elif event.key == pygame.K_1:
                        self.wave_gesture()

            # Walking
            keys = pygame.key.get_pressed()
            x = y = theta = 0.0
            
            # 3. SAFER TOP SPEED: Dropped from 0.9 (sprinting) to 0.6 (safe walking)
            if keys[pygame.K_w] or keys[pygame.K_UP]:    x = 0.6
            if keys[pygame.K_s] or keys[pygame.K_DOWN]:  x = -0.6
            if keys[pygame.K_a] or keys[pygame.K_LEFT]:  theta = 0.5
            if keys[pygame.K_d] or keys[pygame.K_RIGHT]: theta = -0.5

            if abs(x) > 0.1 or abs(theta) > 0.1:
                # Pass the custom carpet gait to the walk engine
                self.motion.moveToward(x, y, theta, carpet_config)
            else:
                self.motion.stopMove()

            self.handle_head_movement()

            # Camera
            frame = self.get_image()
            if frame is not None:
                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                surf = pygame.surfarray.make_surface(rgb.swapaxes(0,1))
                self.screen.blit(surf, (0, 0))
            else:
                self.screen.fill((10, 10, 30))

            pygame.display.flip()
            self.clock.tick(0)

        # Shutdown
        print("🛑 Shutting down...")
        self.motion.stopMove()
        self.motion.rest()
        if self.video and self.video_client:
            try:
                self.video.unsubscribe(self.video_client)
            except:
                pass
        pygame.quit()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ip", type=str, default="127.0.0.1")
    parser.add_argument("--port", type=int, default=9559)
    args = parser.parse_args()

    session = qi.Session()
    try:
        session.connect(f"tcp://{args.ip}:{args.port}")
        print(f"✅ Connected to {args.ip}")
    except Exception as e:
        print(f"❌ Connection failed: {e}")
        sys.exit(1)

    NaoTeleop(session).run()