Error When Rendering Scene Without Contacts: AttributeError in _compress_contact()

[jc-bao]

When I try to render the scene without any contacts, I encounter the following error:

Traceback (most recent call last):
  File "/home/pcy/Research/code/crazyswarm2-adaptive/utils/redis_minimum/mbd_core.py", line 201, in <module>
    main(tyro.cli(Args))
  File "/home/pcy/Research/code/crazyswarm2-adaptive/utils/redis_minimum/mbd_core.py", line 191, in main
    webpage = html.render(env.sys.tree_replace({"opt.timestep": env.dt}), rollout)
  File "/home/pcy/.local/lib/python3.10/site-packages/brax/io/html.py", line 76, in render
    return render_from_json(json.dumps(sys, states), height, colab, base_url)
  File "/home/pcy/.local/lib/python3.10/site-packages/brax/io/json.py", line 184, in dumps
    states = _compress_contact(states)
  File "/home/pcy/.local/lib/python3.10/site-packages/brax/io/json.py", line 107, in _compress_contact
    c = [compress(states.contact, i) for i in range(states.x.pos.shape[0])]
  File "/home/pcy/.local/lib/python3.10/site-packages/brax/io/json.py", line 107, in <listcomp>
    c = [compress(states.contact, i) for i in range(states.x.pos.shape[0])]
  File "/home/pcy/.local/lib/python3.10/site-packages/brax/io/json.py", line 105, in compress
    lambda x: pad(x[contact_mask[i]], n_contact), contact.take(i)
AttributeError: 'Contact' object has no attribute 'take'

I suspect the error originates from the _compress_contact function, where the shape check previously was:

def _compress_contact(states: State) -> State:
  """Reduces the number of contacts based on penetration > 0."""
  if states.contact is None or states.contact.pos.shape[0] == 0:
    return states

However, if the system has no contact with the environment, should it not be states.contact.pos.shape[1] == 0?

After making that modification, my system works fine. I am not entirely sure if this is the correct way to handle it.