Posted 2007-05-02.
Generators are a powerful feature 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.)
I’ve worked for several years on RoboCup, the international robot soccer competition. Our software 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:
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:
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.
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:
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:
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:
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.
#!/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)