Changing geoms on the fly

[yardenas]

Hi!

I'm trying to change the geometry of objects as part of my domain randomization pipeline. For example, I'd like to sample the length of the pole in pendulum environment. Is that possible?

Edit: I'm pretty sure that the following code should work

def domain_randomization(sys, rng):
    @jax.vmap
    def randomize(rng):
        cpole = jax.random.normal(rng) * 0.05 + sys.geom_size[-1, 1]
        length = sys.geom_size.at[-1, 1].set(cpole)
        return length, cpole

    length, samples = randomize(rng)
    in_axes = jax.tree_map(lambda x: None, sys)
    in_axes = in_axes.tree_replace({"geom_size": 0})
    sys = sys.tree_replace({"geom_size": length})
    return sys, in_axes, samples[:, None]

Edit 2: The above code doesn't seem to work. While geom_size does change in sys, for some reason when propagating the dynamics one step ahead, I get exactly the same next state, as if there's not domain randomization at all.