Showing multiple episodes at once

[valentinmace]

Hi, I would like to be able to render multiple episodes at once. Precisely, I have multiple policies trained for a Brax environment (say Ant) and I would like to render all their trajectory in the same video/GIF as in the example below.

What would be the best way to do it ? Thanks.

[btaba]

Hi @valentinmace , one way to do it in new brax (>0.1.0),is to generate xml with unique body names (as many as you need for each trajectory) and pass stacked state.x to the renderer manually

Here's a code snippet where you stack state.x for a list of trajectories in trajectory_states:

stacked_states = [states[0].replace(x=jax.tree_map(lambda *x: jp.concatenate(x), *[s.x for s in states]))
              for states in trajectory_states]