supplement vision randomization function probable error

[alexeiplatzer]

The _supplement_vision_randomization_fn in wrapper handles missing "geom_matid" in a different way from the _identity_vision_randomization_fn. the identity function works as expected, as do custom randomization functions that set the geom_matid themselves. however, custom randomization functions that do not set geom_matid fail with the madrona mjx renderer. It seems that the problem is that the geom_matid is not broadcasted to the correct shape.
In supplement function:

val = -1 if field == 'geom_matid' else getattr(mjx_model, field)
      mjx_model = mjx_model.tree_replace({
          field: jp.repeat(jp.expand_dims(val, 0), num_worlds, axis=0),
      })

Compare to indentity function:

'geom_matid': jp.repeat(
          jp.expand_dims(jp.repeat(-1, mjx_model.geom_matid.shape[0], 0), 0),
          num_worlds,
          axis=0,
      ),

an additional 'repeat' for the value -1 is likely missing, giving the error. as -1 is a scalar value, unlike the geom_matid itself, it probably results in a shape mismatch, which is probably why this supplement function doesn't work with missing geom_matid.
Probably can be very easily fixed by setting val not to -1 but to the same as in identity fucntion.