# Creating robot behaviors with Python generators

*Posted 2007-05-02.*

[Generators](https://en.wikipedia.org/wiki/Generator_(computer_programming)) are a [powerful feature](https://docs.python.org/2.7/reference/simple_stmts.html#the-yield-statement) of the Python programming language. In a nutshell, generators let you write a function that behaves like an iterator. The standard approach to programming robot behaviors is based on state machines. However, robotics code is *full* of special cases, so a complex behavior will typically end up with a lot of bookkeeping cruft. Generators let us simplify the bookkeeping and express the desired behavior in a straightforward manner.

(Idea originally due to [Jim Bruce](http://www.cs.cmu.edu/~jbruce/).)

I've worked for several years on [RoboCup](https://robocup.org), the international robot soccer competition. [Our software](http://www.cs.cmu.edu/~robosoccer/legged/) is written in a mixture of C++ (for low-level localization and vision algorithms) and Python (for high-level behaviors). Let's say we want to write a simple goalkeeper for a robot soccer team. Our keeper will be pretty simple; here's a list of the requirements:

  1. If the ball is far away, stand in place.
  2. If the ball is near by, dive to block it. Dive to the left if the ball is to the left; dive to the right if the ball is to the right.
  3. If we choose a "dive" action, then "stand" on the next frame, nothing will happen. (Well, maybe the robot will twitch briefly....) So when we choose to dive, we need to commit to sending the same dive command for some time (let's say one second).

The usual approach to robot behavior design relies on hierarchical state machines. Specifically, we might be in a "standing" state while the ball is far away; when the ball becomes close, we enter a "diving" state that persists for one second. Because of requirement 3, this solution will have a few warts: we need to keep track of how much time we've spent in the dive state. Every time we add a special case like this, we need to keep some extra state information around. Since robotics code is full of special cases, we tend to end up with a lot of bookkeeping cruft. In contrast, generators will let us clearly express the desired behavior.

On to the state-machine approach. First, we'll have a class called Features that abstracts the robot's raw sensor data. For this example, we only care whether the ball is near/far and left/right, so Features will just contain two boolean variables:

```python
class Features(object):
    ballFar = True
    ballOnLeft = True
```

Next, we make the goalkeeper. The keeper's behavior is specified by the `next()` function, which is called thirty times per second by the robot's main event loop (every time the on-board camera produces a new image). The `next()` function returns one of three actions: `"stand"`, `"diveLeft"`, or `"diveRight"`, based on the current values of the Features object. For now, let's pretend that requirement 3 doesn't exist.

```python
class Goalkeeper(object):
    def __init__(self, features):
        self.features = features

    def next(self):
        features = self.features
        if features.ballFar:
            return 'stand'
        else:
            if features.ballOnLeft:
                return 'diveLeft'
            else:
                return 'diveRight'
```

That was simple enough. The constructor takes in the `Features` object; the `next()` method checks the current `Features` values and returns the correct action. Now, how about satisfying requirement 3? When we choose to dive, we need to keep track of two things: how long we need to stay in the `"dive"` state and which direction we dove. We'll do this by adding a couple of instance variables (`self.diveFramesRemaining` and `self.lastDiveCommand`) to the Goalkeeper class. These variables are set when we initiate the dive. At the top of the `next()` function, we check if `self.diveFramesRemaining` is positive; if so, we can immediately return `self.lastDiveCommand` without consulting the `Features`. Here's the code:

```python
class Goalkeeper(object):
    def __init__(self, features):
        self.features = features
        self.diveFramesRemaining = 0
        self.lastDiveCommand = None

    def next(self):
        features = self.features
        if self.diveFramesRemaining > 0:
            self.diveFramesRemaining -= 1
            return self.lastDiveCommand
        else:
            if features.ballFar:
                return 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                self.lastDiveCommand = command
                self.diveFramesRemaining = 29
                return command
```

 This satisfies all the requirements, but it's ugly. We've added a couple of bookkeeping variables to the Goalkeeper class. Code to properly maintain these variables is sprinkled all over the `next()` function. Even worse, the structure of the code no longer accurately represents the programmer's intent: the top-level if-statement depends on the state of the robot rather than the state of the world. The intent of the original `next()` function is much easier to discern. (In real code, we could use a state-machine class to tidy things up a bit, but the end result would still be ugly when compared to our original `next()` function.)

With generators, we can preserve the form of the original `next()` function and keep the bookkeeping only where it's needed. If you're not familiar with generators, you can think of them as a special kind of function. The `yield` keyword is essentially equivalent to `return`, but the next time the generator is called, *execution continues from the point of the last `yield`*, preserving the state of all local variables. With `yield`, we can use a `for` loop to "return" the same dive command the next 30 times the function is called! Lines 11-16 of the below code show the magic:

```python
class GoalkeeperWithGenerator(object):
    def __init__(self, features):
        self.features = features

    def behavior(self):
        while True:
            features = self.features
            if features.ballFar:
                yield 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                for i in xrange(30):
                    yield command
```

Here's a simple driver script that shows how to use our goalkeepers:

```python
import random

f = Features()
g1 = Goalkeeper(f)
g2 = GoalkeeperWithGenerator(f).behavior()

for i in xrange(10000):
    f.ballFar = random.random() > 0.1
    f.ballOnLeft = random.random() < 0.5
    g1action = g1.next()
    g2action = g2.next()
    print "%s\t%s\t%s\t%s" % (
        f.ballFar, f.ballOnLeft, g1action, g2action)
    assert(g1action == g2action)
```

... and we're done! I hope you'll agree that the generator-based keeper is much easier to understand and maintain than the state-machine-based keeper. You can grab the full source code below and take a look for yourself.

```python
#!/usr/bin/env python

class Features(object):
    ballFar = True
    ballOnLeft = True


class Goalkeeper(object):
    def __init__(self, features):
        self.features = features
        self.diveFramesRemaining = 0
        self.lastDiveCommand = None

    def next(self):
        features = self.features
        if self.diveFramesRemaining:
            self.diveFramesRemaining -= 1
            return self.lastDiveCommand
        else:
            if features.ballFar:
                return 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                self.lastDiveCommand = command
                self.diveFramesRemaining = 29
                return command


class GoalkeeperWithGenerator(object):
    def __init__(self, features):
        self.features = features

    def behavior(self):
        while True:
            features = self.features
            if features.ballFar:
                yield 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                for i in xrange(30):
                    yield command


import random
f = Features()
g1 = Goalkeeper(f)
g2 = GoalkeeperWithGenerator(f).behavior()

for i in xrange(10000):
    f.ballFar = random.random() > 0.1
    f.ballOnLeft = random.random() < 0.5
    g1action = g1.next()
    g2action = g2.next()
    print "%s\t%s\t%s\t%s" % (
        f.ballFar, f.ballOnLeft, g1action, g2action)
    assert(g1action == g2action)
```
