A program for LEGO MINDSTORMS Robot Inventor - Gelo model

gelo.py

# SPDX-License-Identifier: MIT
# Copyright (c) 2021 David Lechner <david@pybricks.com>

# A program for LEGO MINDSTORMS Robot Inventor - Gelo

# Developed using MINDSTORMS App v1.3.4 (10.1.0), hub firmware v1.2.01.0103

# Building instructions: https://www.lego.com/cdn/product-assets/product.bi.additional.main.pdf/51515_Gelo.pdf
# Hub API: https://lego.github.io/MINDSTORMS-Robot-Inventor-hub-API/

import hub
from urandom import randint
from ustruct import pack
from utime import ticks_diff, ticks_ms, sleep_ms

COLOR_SENSOR_ID = const(61)
ULTRASONIC_SENSOR_ID = const(62)

MODE_PWM = const(0)
MODE_SPEED = const(1)
MODE_POS = const(2)
MODE_ABS_POS = const(3)

FORMAT_RAW = const(0)
FORMAT_PCT = const(1)
FORMAT_SI = const(2)

STOP_FLOAT = const(0)
STOP_BRAKE = const(1)
STOP_HOLD = const(2)

BUSY_MODE = const(0)
BUSY_MOTOR = const(1)

EVENT_COMPLETED = const(0)
EVENT_INTERRUPTED = const(1)
EVENT_STALLED = const(2)

# position where foot is flat on the ground relative
# to absolute position marker on the motor
ABS_FLAT = const(-40)

# size of step in degrees for transitioning from
# one leg position to another (should add up to 360)
STEP_FLAT_TO_BACK = const(75)
STEP_BACK_TO_FRONT = const(210)
STEP_FRONT_TO_FLAT = const(75)

# speed while legs are going to/from flat
STEP_SPEED_SLOW = const(50)
# speed while legs are going from back to forward
STEP_SPEED_FAST = const(100)


# from Python standard library itertools
def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
    n = start
    while True:
        yield n
        n += step

# from Python standard library itertools
def repeat(object, times=None):
    # repeat(10, 3) --> 10 10 10
    if times is None:
        while True:
            yield object
    else:
        for _ in range(times):
            yield object

# from Python standard library itertools
def zip_longest(*args, fillvalue=None):
    # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
    iterators = [iter(it) for it in args]
    num_active = len(iterators)

    if not num_active:
        return

    while True:
        values = []
        for i, it in enumerate(iterators):
            try:
                value = next(it)
            except StopIteration:
                num_active -= 1
                if not num_active:
                    return
                iterators[i] = repeat(fillvalue)
                value = fillvalue
            values.append(value)
        yield tuple(values)

# spike.control.Timer doesn't allow decimal points
class Timer:
    """
    Replacement Timer class that allows decimal points so we can measure times of less than one second.
    """
    def __init__(self):
        self.start_ticks = ticks_ms()

    def now_ms(self):
        """Returns the time in milliseconds since the timer was last reset."""
        return ticks_diff(ticks_ms(), self.start_ticks)

    def now(self):
        """Returns the time in seconds since the timer was last reset."""
        return self.now_ms() / 1000

    def reset(self):
        """Resets the timer."""
        self.start_ticks = ticks_ms()

def wait_ms(time):
    """
    Wait for milliseconds
    """
    timer = Timer()
    while timer.now_ms() < time:
        yield



def wait(time):
    """
    Wait for seconds
    """
    timer = Timer()
    while timer.now() < time:
        yield


class Task:
    def __init__(self, generator):
        self._generator = generator()
        self.done = False

    def run_one(self):
        try:
            next(self._generator)
        except StopIteration:
            self.done = True

def run_tasks(*tasks, interval=10):
    """
    Runs tasks
    """
    tasks = [Task(t) for t in tasks]
    timer = Timer()

    while True:
        timer.reset()

        for t in tasks:
            if not t.done:
                t.run_one()

        # sleep for remaining time in interval, if any
        sleep_ms(interval - timer.now_ms())
    

class Motor:
    def __init__(self, port, reverse=False):
        """
        Initalizes a new motor object.

        Args:
            port (str): The name of the port ("A", "B", etc.).
            reverse (bool): If ``True``, reverse the direction of positive rotation.
        """
        self._motor = getattr(hub.port, port).motor

        if self._motor is None:
            raise RuntimeError("Missing motor on port {0}".format(port))

        self._reverse = -1 if reverse else 1
        self._step = 0
        self.reset_position()
        self.reset_step()

    def position(self):
        """
        Gets the current position.

        Returns:
            int: The position in degrees.
        """
        self._motor.mode([(MODE_POS, 0)])
        return self._reverse * self._motor.get()[0]

    def abs_pos(self):
        """
        Gets the current absolute position.

        Returns:
            int: The absolute position in degrees (-180 to 180).
        """
        self._motor.mode([(MODE_ABS_POS, 0)])
        return self._reverse * self._motor.get()[0]

    def set_position(self, position=0):
        """
        Sets the current position to ``position``.

        Args:
            position (int): The new position in degrees.
        """
        self._motor.preset(self._reverse * position)

    def reset_position(self):
        """
        Resets current position to current absolute position.
        """
        self._motor.mode([(MODE_ABS_POS, 0)])
        self._motor.preset(self._motor.get()[0])

    def run_to_position(self, position, speed=50, **kwargs):
        self._motor.run_to_position(self._reverse * position, speed, **kwargs)
        while self._motor.busy(BUSY_MOTOR):
            yield

    def reset_step(self):
        """
        Resets step counter to current position.
        """
        self._step = self.position()

    def run_step(self, step, speed=50, **kwargs):
        """
        Runs the motor relative to the previous step.
        
        Args:
            step (int): The amount to move in degrees.
            speed (int): The max speed of the motor while running in percent.
        """
        self._step += step
        yield from self.run_to_position(self._step, speed, **kwargs)

# LEGO UART I/O device protocol for write_direct() method
# https://github.com/pybricks/technical-info/blob/master/uart-protocol.md
MESSAGE_DATA = const(0xC0)
LENGTH_4 = const(0x10)

class ColorSensor:
    def __init__(self, port):
        """
        Creates a new color sensor connected to ``port``.

        Args:
            port (str): The name of the port ("A", "B", etc.).
        """
        port = getattr(hub.port, port)
        info = port.info()
        if info["type"] != COLOR_SENSOR_ID:
            raise RuntimeError("Expecting color sensor on port {0}".format(port))
        self._dev = port.device
        self._mode_map = {m["name"]: i for i, m in enumerate(info["modes"])}

    def light(self, b1, b2, b3):
        """
        Turn the light on

        Args:
            b1 (int): Brightness of segment 1 (0 to 100).
            b2 (int): Brightness of segment 2 (0 to 100).
            b3 (int): Brightness of segment 3 (0 to 100).
        """
        mode = self._mode_map["LIGHT"]
        self._dev.mode(mode)
        self._dev.write_direct(pack("<5B", MESSAGE_DATA | LENGTH_4 | mode, b1, b2, b3, 0))


class UltrasonicSensor:
    def __init__(self, port):
        """
        Creates a new ultrasonic sensor connected to ``port``.

        Args:
            port (str): The name of the port ("A", "B", etc.).
        """
        self._port = getattr(hub.port, port)
        info = self._port.info()
        if info["type"] != ULTRASONIC_SENSOR_ID:
            raise RuntimeError("Expecting ultrasonic sensor on port {0}".format(port))
        self._dev = self._port.device
        self._mode_map = {m["name"]: i for i, m in enumerate(info["modes"])}

    def light(self, b1, b2, b3, b4):
        """
        Turn the light on

        Args:
            b1 (int): Brightness of segment 1 (0 to 100).
            b2 (int): Brightness of segment 2 (0 to 100).
            b3 (int): Brightness of segment 3 (0 to 100).
            b4 (int): Brightness of segment 4 (0 to 100).
        """
        mode = self._mode_map["LIGHT"]
        self._dev.mode(mode)
        self._dev.write_direct(pack("<5B", MESSAGE_DATA | LENGTH_4 | mode, b1, b2, b3, b4))


class Gelo:
    _rear_right = Motor("A", reverse=True)
    _rear_left = Motor("B")
    _front_right = Motor("C", reverse=True)
    _front_left = Motor("D")
    _ultrasonic = UltrasonicSensor("E")
    _color = ColorSensor("F")

    def _motors(self):
        yield self._rear_right
        yield self._rear_left
        yield self._front_right
        yield self._front_left

    def reset(self, position=ABS_FLAT):
        """
        Resets legs back to ``position``.
        """
        for m in self._motors():
            m.reset_position()

        yield from zip_longest(
            self._rear_right.run_to_position(position),
            self._rear_left.run_to_position(position),
            self._front_right.run_to_position(position),
            self._front_left.run_to_position(position),
        )

    def random(self):
        """
        Sets legs to random position (useful for testing reset).
        """
        for m in self._motors():
            m.reset_position()

        yield from zip_longest(
            self._rear_right.run_to_position(randint(-180, 180)),
            self._rear_left.run_to_position(randint(-180, 180)),
            self._front_right.run_to_position(randint(-180, 180)),
            self._front_left.run_to_position(randint(-180, 180)),
        )

    def cycle_lights(self, delay=0.25):
        while True:
            self._color.light(5, 5, 5)
            self._ultrasonic.light(0, 0, 0, 0)
            yield from wait(delay)
            self._color.light(0, 0, 0)
            self._ultrasonic.light(20, 0, 20, 0)
            yield from wait(delay)
            self._color.light(0, 0, 0)
            self._ultrasonic.light(0, 20, 0, 20)
            yield from wait(delay)

    def walk(self):
        # legs start in "flat" position
        yield from self.reset()

        # flat position is defined as 0
        for m in self._motors():
            m.set_position(0)
            m.reset_step()

        move_left = self._front_left
        move_right = self._rear_right
        flat_left = self._rear_left
        flat_right = self._front_right

        # move two opposite legs to back position
        yield from zip_longest(
            move_left.run_step(STEP_FLAT_TO_BACK, STEP_SPEED_SLOW, stop=STOP_FLOAT),
            move_right.run_step(STEP_FLAT_TO_BACK, STEP_SPEED_SLOW, stop=STOP_FLOAT),
        )

        # repeat walking cycle forever
        while True:
            # move back legs to front, keep flat legs flat
            yield from zip_longest(
                move_left.run_step(STEP_BACK_TO_FRONT, STEP_SPEED_FAST, stop=STOP_FLOAT),
                move_right.run_step(STEP_BACK_TO_FRONT, STEP_SPEED_FAST, stop=STOP_FLOAT),
            )

            # move forward legs to flat and flat legs to back at same time
            yield from zip_longest(
                move_left.run_step(STEP_FRONT_TO_FLAT, STEP_SPEED_SLOW, stop=STOP_FLOAT),
                move_right.run_step(STEP_FRONT_TO_FLAT, STEP_SPEED_SLOW, stop=STOP_FLOAT),
                flat_left.run_step(STEP_FLAT_TO_BACK, STEP_SPEED_SLOW, stop=STOP_FLOAT),
                flat_right.run_step(STEP_FLAT_TO_BACK, STEP_SPEED_SLOW, stop=STOP_FLOAT),
            )

            # swap moving and flat legs
            flat_left, move_left = move_left, flat_left
            flat_right, move_right = move_right, flat_right



# Create your objects here.

gelo = Gelo()



# Write your tasks here.

def test():
    yield from zip(gelo.walk(), gelo.cycle_lights(), wait(30))
    raise SystemExit



# main program

run_tasks(test)