Matplotlib upgrade

[zombie-einstein]

Hey @sash-a, thought I'd grab a look at #253.

Changes here update Matplotlib version and then makes changes to pass type checking, so it builds and passes tests.

But ...

The new API now requires that the animation update function returns a list of artists that have been updated, these are objects returned from plotting calls like a plot object, or text box (see the table here). The state/data of these objects can be updated without redrawing the whole thing.

I've done this rubiks cube environment:

# Create the first frame of the plot and return a list of drawn image artists
images = self._draw(ax, states[0])

def make_frame(state: State) -> Sequence[Artist]:
    # Update the data in each image from the state
    for i, image in enumerate(images):
        image.set_data(state.cube[i])
    # Return the updated images
    return images

# Create the animation object.
self._animation = matplotlib.animation.FuncAnimation(
    fig,
    make_frame,
    frames=states,
    interval=interval,
)

Doing this for all environments is easier for some than others given how they've been written, it means passing all the relevant artists to the animate function when the plot is created, and then correctly updating them inside the update function.

I think the animations will still work as is, as they still generally update the axes by reference inside the animation function, it's just less efficient then the new API that just updates the relevant artists.

[sash-a]

Wow this is an awesome addition, thanks @zombie-einstein!

Which environments are still remaining?

Also I think the way you've done it with make_frame looks good, but I have no experience with the new API

[zombie-einstein]

Which environments are still remaining?

Apart from rubix cube they all need a look at, with varying degrees of complexity. I think the general pattern needs to be something like a draw function that draws the first frame (and also draws the individual frames via the render method), and returns any actors that will be updated during the animation, then use this to initialise the first frame and update actors inside the update-func.

Some also currently have actors that are created dynamically during plotting I've been trying to work out the best pattern for, in some cases it seems it would be more efficient to initialise all the actors ahead of time, and update their states as needed. E.g. for the sudoko environment initialise all the text actors and update the text during animation (rather than create them on the fly).

So there's a decent chunk of work to refactor all the plots. One option could be merge something like this PR, which will updated the requirement and passes linting/tests, and still create plots (just less efficiently), and then incrementally update the individual viewers in separate PRs to make best use of the new API.

[sash-a]

Hi @zombie-einstein sorry for leaving this hanging, I needed a bit of time off 😂

I'll have a deeper look at it this week and give some a proper opinion 😄

[zombie-einstein]

Hey @sash-a I fixed the graph-colouring, MMST, sodoku, and TSP viewers (still need to work out the img path for the knapsack viewer).

Is the animate function intended to display states over multiple episodes? For example for the MMST env, the structure of the graph is resampled when the environment is reset. It's nice if the edges were static, then you don't need to recalculate the node positions etc and can just update the edge & node properties, and also save on the relatively node position calculation. But then this causes an issue if the states span multiple episodes. There's maybe a way to do this where you check if the episode has terminated, though a bit more involved.

[sash-a]

So it currently does work over multiple episodes and I think we should try keep that functionality. Maybe there is a way to tell if an episode is done from the state?

[sash-a]

For knapsack what should work is if you follow the method used in TSP for loading images. I think all that should be needed to move the file into the jumanji repo instead of reading it from docs as the manifest automatically includes png files which the knapsack image is

We recently fixed TSP and I assume missed knapsack.

[zombie-einstein]

So it currently does work over multiple episodes and I think we should try keep that functionality. Maybe there is a way to tell if an episode is done from the state?

Maybe comparing what has changes between two steps? Though would be very environment specific. Some more generic solutions could be

• To include the episode completion flag with the sequence of state, but this is quite a big API change. You;d then have something like reset and step methods for the viewer that redraw or just modify artists.
• Lower impact could be to initialise the required artists, but then redraw them all each step. E.g. for graph edges initialise all the line elements, but then update their start-end co-ords every step. As long as number of edges (i.e. the number of artists) is constant. You still have the impact of doing something like the node position calculations every step.

For knapsack what should work is if you follow the method used in TSP for loading images. I think all that should be needed to move the file into the jumanji repo instead of reading it from docs as the manifest automatically includes png files which the knapsack image is

Cheers that's an easy enough fix then

[sash-a]

I would go with the lower impact option as the other option is a big API change. I don't think it's a big issue if the rendering is a bit slower than it could be 🤷

[zombie-einstein]

Hey @sash-a, I think this is all working now, though some of the code definitely needs a work over. Turned out to be a bit more involved:

• I went with checking for a new episode from changes in state, then updating the static elements for a new episode
• I worked out how to add and remove artists inside the update function, so used this in the cases where the number of artists changes, e.g. routes in tsp
• In some cases there was an issue where old elements of the plot were not be in removed after a new episode. I think I've addressed these and the plots now all have the expected behaviour.
• I've not fully updated all the animations, in the case where they seem to work using the old functionality I've just re-used that.

I also grabbed a couple other changes whilst looking at this:

• I believe use of pkg_resources for loading images is deprecated, so updated that
• Pulled out graph layout functionality into a shared module
• Added a script (based on your notebook) to generate animations for all the environments to make testing easier

Wanted to check:

• Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations
• Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

[sash-a]

Hi @zombie-einstein thanks for this I should have time to take a look today

[sash-a]

Hi @zombie-einstein sorry for not having a look at this, the last two weeks have been a bit crazy

Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations

I think that would be nice, but also happy to do this in a subsequent PR?

Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

The main reason I'd like to do this is so that we don't pin to an old matplotlib version anymore so I do think we should update it in all cases, but if it's a bit slower than the old functionality I wouldn't worry because the viewers don't need to be too fast

[sash-a]

This looks good to me. Just a few very minor changes. Can you please explain why there are so many changes for the graph envs? I get there needs to be changes on reset, but not sure why things changed outside of the if not jnp.array_equal(...) block? (Not saying that it's wrong just curious why)

[sash-a]

Can you rename the file to graph_view_utils

[sash-a]

Since this is in examples can you rather call this visualize_random_agent.py

[sash-a]

Why did this need to change?

[zombie-einstein]

So (somewhat annoyingly) if you draw nodes/points with a circle-artist the scale (if I remember rightly) is a fraction of the region (hence the fractional values), but using a scatter plot the larger integer value gives sensible node sizes.
This could change though, I think it's generally sensible to draw all the points in one go using a scatter, rather than individual artists.

[sash-a]

Can you please explain in a comment what's happening here I'm not quite following

[zombie-einstein]

Yeah will do, I may add a longer comment on the interface outlining how some of the animation functionality is implemented.

In this case the members of the routes, and number of nodes are not static, and so the artists representing them need to be fully replaced, by removing the current ones. You can't pass state (I think) as part of the argument to the update function and so need to keep references to the existing artists in the scope outside the update function. So in the above the steps are:

• Pop the current artists from the outer list
• Call the remove method to remove them from the plot
• Add them to the list of updated artists returned from the update function
• Create new artists and push them in the outer list for use in later updates

I've used a list to act as mutable storage for artists mainly for convenience, maybe a global or class variable would look nicer. Also quite a few lines of code here, was looking to review if there is a neater way to implement.

[zombie-einstein]

This looks good to me. Just a few very minor changes. Can you please explain why there are so many changes for the graph envs? I get there needs to be changes on reset, but not sure why things changed outside of the if not jnp.array_equal(...) block? (Not saying that it's wrong just curious why)

Sure, so a couple things for graphs in general:

• For the graphs the reset generally redraws the full plot, and in some cases relatively heavy computation to layout the nodes (which is REALLY slow if done every step). So the full redraw is only done on a new episode, and means removing all the existing artists and replacing them with the updated ones using some mutable storage for artists outside the animation update function.
• In some cases once the nodes and edges are drawn you can just update their properties inside an episode, e.g. updating their colour without needing to redraw them.
• In the worse cases the plot has features that vary in size every step, e.g. in crvp the route changes every step, so need to update all the artists, and so for convenience I've just been using the existing functionality to remove the existing artists and replace them with newly created ones.

Compared to other environments the graphs are a pain due to all the nodes/edges moving episode, and the potential variable number of artists each step. In comparison for something like Suduko the grid is static, and you know there is a fixed grid of numbers that can just be updated in the animation, you don't need to mess about adding and removing artists inside the update function!

[zombie-einstein]

Hi @zombie-einstein sorry for not having a look at this, the last two weeks have been a bit crazy

Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations

I think that would be nice, but also happy to do this in a subsequent PR?

Cool, I might just grab it, as where implementations vary it kind of makes it hared to refactor when trying to work out what different implementations are doing.

Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

The main reason I'd like to do this is so that we don't pin to an old matplotlib version anymore so I do think we should update it in all cases, but if it's a bit slower than the old functionality I wouldn't worry because the viewers don't need to be too fast

Everything is functional now with the updated version (the graphs were really the ones that didn't play nicely with the new API) updates would be mostly about efficient use of new API, e.g. updating artists rather then continually redrawing. My main concern was a the size of the PR, but I can chip away at the other envs if you don't mind it.