Changing HField geometry at runtime - visualization not working

[kylestach]

I'd like to change the geometry of a HField at runtime. I am using the following code:

hfield = physics.bind(self._hfield_chunks[geom_idx])
res = hfield.nrow

# Generate terrain in a 2D numpy array bounded to [0, 1]...

start_idx = hfield.adr
physics.model.hfield_data[start_idx:start_idx + res**2] = terrain.ravel()

with physics.contexts.gl.make_current() as ctx:
    ctx.call(mjlib.mjr_uploadHField,
             physics.model.ptr,
             physics.contexts.mujoco.ptr,
             hfield.element_id)

Unfortunately, while the physics model seems to change as expected, the visualization is not updating to match. In case it's relevant, I'm using dm_control and the above code is called by the after_step() callback.