Coroutines for sequential behaviors

[jpreiss]

A long time ago, we discussed using coroutines for sequential behaviors, but it required a lot of boilerplate code in Python 2.7. With Python 3 the yield from keyword fixes the problem.

Here is a toy example. We first need a few library functions. The "payoff" is in the complete behavior takeoff_circle_land which would be somewhat cumbersome to implement as a state machine. Note how easy it is to have one control policy sleeping while executing a high level command while others are sending low-level setpoints in a fast loop.

I am leaving this in discussion for now because the example is not complicated enough to stress the library design. If anyone validates/fixes the design in a more complex project, we can promote it to an example script.

"""EXPERIMENTAL: Example of coroutines as an alternative to state machines.

We use the term 'control coroutine' , abbreviated 'coco', for a Python
generator whose inputs are floating-point time values. These generators can be
used to implement sequential behaviors with internal state that might become
cumbersome for a a state-machine-based implementation.
"""

import numpy as np
from pycrazyswarm import *


def begin_at_zero(coco_maker):
    """Decorator for shifting the time argument of a control coroutine.

    Inside the control coroutine, the first value of `yield` will be zero, no
    matter the current real time.
    """
    def coco_wrapper(*args, **kwargs):
        coco = coco_maker(*args, **kwargs)
        coco.send(None)
        t = yield
        t0 = t
        try:
            while True:
                t = yield coco.send(t - t0)
        except StopIteration:
            return
    return coco_wrapper


@begin_at_zero
def ellipse(cf, center, major, minor, period):
    """Infinite (never stops yielding) control coroutine for ellipse."""

    # Scale factors.
    s = 2 * np.pi / period
    s2 = s * s

    while True:
        t = yield
        cos = np.cos(s * t)
        sin = np.sin(s * t)

        pos = (cos * major) + (sin * minor) + center
        vel = -(s * sin * major) + (s * cos * minor)
        acc = -(s2 * cos * major) - (s2 * sin * minor)

        # TODO: Compute the proper value of omega by extracting the needed math
        # into a standalone function in uav_trajectory.py.
        omega = np.zeros(3)

        cf.cmdFullState(
            pos,
            vel,
            acc,
            yaw=0.0,
            omega=omega,
        )


@begin_at_zero
def truncate(coco, duration):
    """Executes an infinite control coroutine for a finite duration."""

    coco.send(None)
    t = yield
    while t <= duration:
        coco.send(t)
        t = yield


@begin_at_zero
def sleep(duration):
    """Coroutine version of timeHelper.sleep() for high-level commands."""

    t = yield
    while t <= duration:
        t = yield


def poll_cocos(cocos, timeHelper, rateHz):
    """Executes a set of control coroutines."""

    for coco in cocos:
        coco.send(None)

    while True:
        t = timeHelper.time()
        n_done = 0
        for coco in cocos:
            try:
                coco.send(t)
            except StopIteration:
                n_done += 1
        if n_done == len(cocos):
            break
        timeHelper.sleep(1.0/rateHz)


def takeoff_circle_land(cf, loops):
    """Example of a sequential behavior."""

    print(f"cf{cf.id} taking off.")
    cf.takeoff(1.0, duration=3.0)
    yield from sleep(4.0)

    print(f"cf{cf.id} beginning circle.")
    major = np.array([1.0, 0.0, 0.0])
    minor = np.array([0.0, 1.0, 0.0])
    center = cf.position() - major
    period = 12.0
    yield from truncate(
        ellipse(cf, center, major, minor, period),
        duration=loops*period
    )

    print(f"cf{cf.id} landing.")
    cf.notifySetpointsStop()
    cf.land(0.05, duration=3.0)
    yield from sleep(4.0)


def main():
    crazyflies_yaml = """
    crazyflies:
    - channel: 100
      id: 1
      initialPosition: [1.0, 0.0, 0.0]
    - channel: 100
      id: 2
      initialPosition: [0.0, -1.0, 0.0]
    """

    swarm = Crazyswarm(crazyflies_yaml=crazyflies_yaml)
    timeHelper = swarm.timeHelper
    allcfs = swarm.allcfs
    cfs = allcfs.crazyflies

    cocos = [
        takeoff_circle_land(cf, loops=i+1)
        for i, cf in enumerate(cfs)
    ]
    poll_cocos(cocos, timeHelper, rateHz=30)


if __name__ == "__main__":
    main()