Quiero agregar una leyenda en una animación de Python, como elline.set_label() abajo. Esto es similar aplt.plot(x,y,label='%d' %*variable*).

Sin embargo, encuentro que los códigos no funcionan aquí. La animación solo muestra líneas que cambian pero no hay etiquetas o leyendas disponibles. ¿Como puedo solucionar este problema?

 from matplotlib import pyplot as plt 
 from matplotlib import animation

fig = plt.figure()

ax = plt.axes(xlim=(0, 2), ylim=(0, 100))

N = 3
lines = [plt.plot([], [])[0] for _ in range(N)]


def init():    
    for line in lines:
        line.set_data([], [])
    return lines

def animate(i):
    for j,line in enumerate(lines):
        line.set_data([0, 2], [10*j,i])
        line.set_label('line %d, stage %d'%(j,i))
    return lines

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=100, interval=20, blit=True)

plt.show()