Arm simulation visualization with Matplotlib

One of the downsides of switching to Python from Matlab is that it can be a pain to plot some kinds of things, and I’ve found animations to be one those things. In previous posts I’ve done the visualization of my arm simulations through Pyglet, but I recently started playing around with Matplotlib’s animation function, and the results are pretty smooth. The process is also relatively painless and quick to get up and running, so I thought I would throw up some Matplotlib code for visualizing my previously described 2 link arm MapleSim simulation.

So, let’s look at the code:

#Written by Travis DeWolf (Sept, 2013)
#Based on code by Jake Vanderplas - http://jakevdp.github.com

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import py2LinkArm

class TwoLinkArm:
    """
    :param list u: the torque applied to each joints
    """
    def __init__(self, u = [.1, 0]): 
        self.u = np.asarray(u, dtype='float') # control signal
        self.state = np.zeros(3) # vector for current state
        self.L1=0.37 # length of arm link 1 in m
        self.L2=0.27 # length of arm link 2 in m
        self.time_elapsed = 0

        self.sim = py2LinkArm.pySim()
        self.sim.reset(self.state)
    
    def position(self):
        """Compute x,y position of the hand"""

        x = np.cumsum([0,
                       self.L1 * np.cos(self.state[1]),
                       self.L2 * np.cos(self.state[2])])
        y = np.cumsum([0,
                       self.L1 * np.sin(self.state[1]),
                       self.L2 * np.sin(self.state[2])])
        return (x, y)

    def step(self, dt):
        """Simulate the system and update the state"""
        for i in range(1500):
            self.sim.step(self.state, self.u)
        self.time_elapsed = self.state[0]

#------------------------------------------------------------
# set up initial state and global variables
arm = TwoLinkArm()
dt = 1./30 # 30 fps

#------------------------------------------------------------
# set up figure and animation
fig = plt.figure(figsize=(4,4))
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                     xlim=(-1, 1), ylim=(-1, 1))
ax.grid()

line, = ax.plot([], [], 'o-', lw=4, mew=5)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def init():
    """initialize animation"""
    line.set_data([], [])
    time_text.set_text('')
    return line, time_text

def animate(i):
    """perform animation step"""
    global arm, dt
    arm.step(dt)
    
    line.set_data(*arm.position())
    time_text.set_text('time = %.2f' % arm.time_elapsed)
    return line, time_text

# frames=None for matplotlib 1.3
ani = animation.FuncAnimation(fig, animate, frames=None,
                              interval=25, blit=True, 
                              init_func=init)

# uncomment the following line to save the video in mp4 format.
# requires either mencoder or ffmpeg to be installed
#ani.save('2linkarm.mp4', fps=15, 
#         extra_args=['-vcodec', 'libx264'])

plt.show()

There’s not all too much to it, which is nice. I’ve created a class TwoLinkArm that wraps the actual arm simulator, stores the current state of the arm, and has a step function that gets called to move the system forward in time. Then, I created a line and stored a reference to it; this is what is going to be updated in the simulation. I then need functions that specify how the line will be initialized and updated. The init function doesn’t do anything except set the line and text data to nothing, and then the animate function calls the arm simulation’s step function and updates the line and text data.

For more details about it there’s this blog post which steps through the process a bit more. For simple arm simulations the above is all I need though, so I’ll leave it there for now!

Here’s an animation of the resulting sim visualization, I’ve removed a lot of the frames to keep the size down. It’s smoother when running the actual code, which can be found up on my github.

Hurray better visualization!

Tagged , , ,

8 thoughts on “Arm simulation visualization with Matplotlib

  1. yizheng says:

    How do I incorporate this into anaconda2 ?

    • travisdewolf says:

      Hi! Should hopefully just need to have a recent version of matplotlib and be good to go, is it not running for you? There might be a few other dependencies you need to pip install…

      • yizheng says:

        Hi ! Thx for the reply ! I mean where should I put all those library files, I tried to put the 2LinkArm folder into ~/anaconda2/lib/python2.7/ but it gives me an import error, says “no module named py2LinkArm”.

        Also, in your previous post about wrapping the MapleSIm model for python, when I ran:
        g++ c2LinkArm_run.cpp -o sim
        ./sim
        It gave me a segmentation fault, and the ‘output.dat’ has nothing in it. Do I only need the two library files in order to run it or do I actually need to purchase the MapleSim software ?

        Thx

      • travisdewolf says:

        Oh, I just keep everything in the same folder and run from there!

        And thanks for catching that seg fault! Not sure how I missed it :S
        I’ve fixed the error (needed to instantiate the out array), if you do another pull and check it out hopefully it should run for you!

      • yizheng says:

        Now it works, but somehow on the PyCharm IDE scroll panel there is still a notification says “no module named py2LinkArm”, but it actually works. Also, for any shared object files, do I just need to put them into the current python project dir (which I believe is also in sys.path) or do I have to reset my LD_LIBRARY_PATH ?

      • travisdewolf says:

        glad to hear! and yup I usually just throw things into the same directory for quick projects.

      • yizheng says:

        Also, I’m actually learning your implementation of DMP, and it helps me a lot ! But it looks like I can set the running time for the canonical system (cs.run_time), I assume this decides the duration of the trajectory when tau=1 right ? So when I set it to 5.0 or 3.0 and run the discrete dmp, there will be an overflow. Another thing is that the when I animate it, the run_time in the program doesn’t really correspond to the actual duration in reality. Any idea how to fix those ?

        Thx!

      • travisdewolf says:

        Hello! Right, so the parameter names aren’t the best. To speed up the run relative to some external system is to set the dt parameter on the DMP. By default it’s .001, so if you increase it your run will go faster, and if you decrease it it will take longer. And to have it correspond to other times you should use the same time step of your plotting animation. There can be issues with stability if the dt parameter is set to high, so what you can do instead is set dt to some stable value lower than your desired time step, and then every time you call the DMP step, call it int(dt_desired / dt) times. Hope that helps!

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: