fix(queries,tasks): return per-env (num_envs,) shape for contact query and draw _terminated (#781)

ContactData raw-MuJoCo branch rebound the (num_envs,) array to a bare `True` on
a detected contact, returning a 0-d tensor while the MJX branch returns
(num_env,) — same query, divergent shape. Write `has_contact[0] = True` (raw
MuJoCo is single-env; MujocoHandler rejects num_envs>1) and return via
from_numpy to keep the (num_envs,) shape.

draw_svg/draw_triangle _terminated returned torch.tensor([False]) — fixed shape
(1,) on the default device — mismatching the (num_envs,) done/timeout
bookkeeping for num_envs>1 and crashing on GPU (CPU terminated vs CUDA timeout).
Return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device),
matching the base _terminated contract.

Adds general regression tests (stub mujoco handler for the query; __new__ tasks
for _terminated). Red on pre-fix.

Author: geng-haoran

roboverse_pack/queries/contact.py

@@ -68,9 +68,12 @@ def __call__(self):
                     con.geom1 == self._geom2_id and con.geom2 == self._geom1_id
                 ):
                     if con.dist < 1e-6:
-                        has_contact = True
+                        has_contact[0] = True
                         break
 
-            return torch.tensor(has_contact, dtype=torch.float32, device=self.handler.device)
+            # Keep the per-env array shape (num_envs,) to match the MJX branch.
+            # Rebinding ``has_contact`` to a bare ``True`` produced a 0-d tensor,
+            # diverging from MJX's (num_env,) output for the same query.
+            return torch.from_numpy(has_contact).to(dtype=torch.float32, device=self.handler.device)
 
         raise ValueError(f"Unsupported handler type: {type(self.handler)} for ContactData query")

roboverse_pack/tasks/maniskill/draw_svg.py

@@ -24,7 +24,7 @@ class DrawSvgTask(ManiskillBaseTask):
     # rewrite terminate
     def _terminated(self, states: TensorState) -> torch.Tensor:
         """No terminate condition yet. Will terminate when time is up."""
-        return torch.tensor([False])
+        return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
 
     # rewrite checker
     def reset(self, states=None, env_ids=None):

roboverse_pack/tasks/maniskill/draw_triangle.py

@@ -24,7 +24,7 @@ class DrawTriangleTask(ManiskillBaseTask):
     # rewrite terminate
     def _terminated(self, states: TensorState) -> torch.Tensor:
         """No terminate condition yet. Will terminate when time is up."""
-        return torch.tensor([False])
+        return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
 
     # rewrite checker
     def reset(self, states=None, env_ids=None):

tests/test_contact_query_shape.py

@@ -0,0 +1,53 @@
+"""Regression: ContactData returns a per-env (num_envs,) tensor on raw MuJoCo.
+
+The raw-MuJoCo branch built ``has_contact = np.zeros(num_envs)`` but, on a
+detected contact, rebound the name to a bare ``True`` — so ``torch.tensor(True)``
+returned a 0-d scalar, while the MJX branch returns shape ``(num_env,)``. Same
+query, divergent output shape across backends. This pins the (num_envs,) shape
+in both the contact-present and contact-absent cases, with a stub handler (no
+mujoco/jax/GPU needed).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+def _mujoco_stub(contacts):
+    """A handler whose module looks like the raw-MuJoCo backend."""
+    stub_cls = type("StubMujocoHandler", (), {})
+    stub_cls.__module__ = "metasim.sim.mujoco.stub"  # selects the raw-MuJoCo branch
+    h = stub_cls()
+    h.num_envs = 1
+    h.device = "cpu"
+    h.data = SimpleNamespace(ncon=len(contacts), contact=contacts)
+    return h
+
+
+def _query(handler):
+    from roboverse_pack.queries.contact import ContactData
+
+    q = ContactData("geom_a", "geom_b")
+    q._geom1_id = 10
+    q._geom2_id = 20
+    q.handler = handler  # bypass bind_handler (which needs a real mj model)
+    return q
+
+
+@pytest.mark.general
+def test_contact_present_returns_per_env_tensor():
+    contacts = [SimpleNamespace(geom1=10, geom2=20, dist=0.0)]  # matching pair, in contact
+    out = _query(_mujoco_stub(contacts))()
+    assert out.shape == (1,), f"expected (num_envs,) = (1,), got {tuple(out.shape)}"
+    assert out.dtype == torch.float32
+    assert out[0] == 1.0
+
+
+@pytest.mark.general
+def test_contact_absent_returns_per_env_tensor():
+    out = _query(_mujoco_stub([]))()  # no contacts
+    assert out.shape == (1,)
+    assert out[0] == 0.0

tests/test_draw_terminated_shape.py

@@ -0,0 +1,37 @@
+"""Regression: draw_svg / draw_triangle _terminated returns a per-env (num_envs,) tensor.
+
+Both tasks returned ``torch.tensor([False])`` — a fixed shape ``(1,)`` on the
+default device — regardless of num_envs, mismatching the ``(num_envs,)`` timeout
+/ done bookkeeping when num_envs > 1. They now return
+``torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)``. Backend-free:
+``_terminated`` ignores its ``states`` argument, so it can be called on a
+``__new__`` instance with num_envs/device set.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+
+@pytest.mark.general
+@pytest.mark.parametrize(
+    "module_name,class_name",
+    [
+        ("roboverse_pack.tasks.maniskill.draw_svg", "DrawSvgTask"),
+        ("roboverse_pack.tasks.maniskill.draw_triangle", "DrawTriangleTask"),
+    ],
+)
+def test_draw_terminated_is_per_env(module_name, class_name):
+    import importlib
+
+    cls = getattr(importlib.import_module(module_name), class_name)
+    task = cls.__new__(cls)  # bypass __init__/backend
+    task.num_envs = 4
+    task.device = torch.device("cpu")
+
+    out = task._terminated(None)  # _terminated ignores states
+
+    assert out.shape == (4,), f"{class_name}._terminated must be (num_envs,), got {tuple(out.shape)}"
+    assert out.dtype == torch.bool
+    assert not out.any()