diff --git a/CHANGELOG.md b/CHANGELOG.md index 510c6978..72ca1864 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,27 @@ All notable changes to GeneLab are recorded here. ## [Unreleased] +### Fixed + +- **Root-velocity writes were silent no-ops on rigid articulations** (#242): Genesis's + `RigidEntity` has `get_vel` / `get_ang` but no `set_vel` / `set_ang` (those setters exist + only on FEM / tool entities), so `push_by_setting_velocity`, the velocity half of + `reset_root_state_uniform`, and `Articulation.write_root_state` wrote nothing — push-based + domain randomization never reached the simulator. Root velocity now routes through the new + `genelab.entity.write_root_velocity`, which drives `set_dofs_velocity` on the base free + joint's 6 DoFs (0–2 world-frame linear, 3–5 angular) and falls back to direct setters for + entity types that have them. The wuji reorient cube kick / reset and the Franka cube reset + in `examples/` had the same dead idiom and now use the shared helper. Policies previously + trained with push events effectively had those disturbances disabled; retraining may be + needed where push robustness matters. +- **Teleop HUD text never displayed under Genesis 1.2** (same bug class as #242, found by + auditing every `getattr`-guarded Genesis call): the outer `Viewer` forwards + `register_keybinds` but not `set_message_text`, so the keyboard bridge's HUD write + silently no-opped. The bridge now falls back to the inner pyrender viewer. A new + `tests/test_genesis_api_contract.py` pins every `getattr`-guarded Genesis method name + against the installed Genesis so future upstream renames fail tests instead of silently + disabling physics or UI. + ## [0.4.0] — 2026-07-03 ### Changed diff --git a/examples/franka/src/genelab_franka/mdp.py b/examples/franka/src/genelab_franka/mdp.py index 46d521e8..5ba2609b 100644 --- a/examples/franka/src/genelab_franka/mdp.py +++ b/examples/franka/src/genelab_franka/mdp.py @@ -11,6 +11,7 @@ import torch +from genelab.entity import write_root_velocity from genelab.entity._torch import to_tensor from genelab_franka.constants import ( @@ -171,14 +172,7 @@ def reset_cube_uniform( except TypeError: set_pos(new_pos) zeros = torch.zeros(n, 3, device=env.device) - for fn_name in ("set_vel", "set_ang"): - fn = getattr(handle, fn_name, None) - if fn is None: - continue - try: - fn(zeros, envs_idx=env_ids) - except TypeError: - fn(zeros) + write_root_velocity(handle, zeros, zeros, env_ids) def resample_goal_uniform( diff --git a/examples/wuji/src/genelab_wuji/reorient/mdp/commands.py b/examples/wuji/src/genelab_wuji/reorient/mdp/commands.py index e94a1696..4516b97a 100644 --- a/examples/wuji/src/genelab_wuji/reorient/mdp/commands.py +++ b/examples/wuji/src/genelab_wuji/reorient/mdp/commands.py @@ -14,6 +14,7 @@ import torch +from genelab.entity import write_root_velocity from genelab.managers.command_manager import CommandTerm, CommandTermCfg from genelab.utils.math import quat_error_magnitude, quat_mul @@ -217,12 +218,11 @@ def _pose_goal_marker(self) -> None: for setter, value in ( ("set_pos", pos), ("set_quat", self._goal_quat_w), - ("set_vel", zeros), - ("set_ang", zeros), ): fn = getattr(handle, setter, None) if fn is not None: fn(value) + write_root_velocity(handle, zeros, zeros) @property def success_achieved(self) -> torch.Tensor: diff --git a/examples/wuji/src/genelab_wuji/reorient/mdp/events.py b/examples/wuji/src/genelab_wuji/reorient/mdp/events.py index f748c9ec..cc153e31 100644 --- a/examples/wuji/src/genelab_wuji/reorient/mdp/events.py +++ b/examples/wuji/src/genelab_wuji/reorient/mdp/events.py @@ -4,6 +4,7 @@ import torch +from genelab.entity import write_root_velocity from genelab_wuji.reorient.constants import REORIENT_CUBE_INIT_POS from genelab_wuji.reorient.mdp._state import cage_counter, disturbance_scale_value from genelab_wuji.reorient.mdp._math import random_quat @@ -31,8 +32,6 @@ def reset_object_orientation( for setter, value in ( ("set_pos", pos), ("set_quat", quat), - ("set_vel", zeros), - ("set_ang", zeros), ): fn = getattr(handle, setter, None) if fn is None: @@ -41,6 +40,7 @@ def reset_object_orientation( fn(value, envs_idx=env_ids) except TypeError: fn(value) + write_root_velocity(handle, zeros, zeros, env_ids) def reset_cage_state(env: "EnvContext", env_ids: torch.Tensor | None) -> None: @@ -89,32 +89,27 @@ def randomize_cube_physics( pass -def _kick( +def _kicked( handle: object, getter: str, - setter: str, env: "EnvContext", env_ids: torch.Tensor, n: int, lo: float, hi: float, scale: float, -) -> None: +) -> torch.Tensor | None: + """Current velocity (via ``getter``) plus a random-direction impulse, for ``env_ids``.""" fn_get = getattr(handle, getter, None) - fn_set = getattr(handle, setter, None) - if fn_get is None or fn_set is None: - return + if fn_get is None: + return None cur = torch.as_tensor(fn_get(), device=env.device, dtype=torch.float) if cur.dim() == 1: cur = cur.unsqueeze(0).expand(env.num_envs, -1) direction = torch.randn(n, 3, device=env.device) direction = direction / direction.norm(dim=-1, keepdim=True).clamp_min(1e-6) mag = torch.empty(n, 1, device=env.device).uniform_(lo, lo + (hi - lo) * scale) - new = cur[env_ids] + direction * mag - try: - fn_set(new, envs_idx=env_ids) - except TypeError: - fn_set(new) + return cur[env_ids] + direction * mag def apply_velocity_disturbance( @@ -137,5 +132,8 @@ def apply_velocity_disturbance( n = int(env_ids.numel()) handle = env.scene[object_name].gs_handle # type: ignore[index] scale = float(disturbance_scale_value(env)[0]) - _kick(handle, "get_vel", "set_vel", env, env_ids, n, min_speed, max_speed, scale) - _kick(handle, "get_ang", "set_ang", env, env_ids, n, min_angular, max_angular, scale) + vel = _kicked(handle, "get_vel", env, env_ids, n, min_speed, max_speed, scale) + ang = _kicked(handle, "get_ang", env, env_ids, n, min_angular, max_angular, scale) + if vel is None or ang is None: + return + write_root_velocity(handle, vel, ang, env_ids) diff --git a/src/genelab/bridges/keyboard.py b/src/genelab/bridges/keyboard.py index d84b2e73..5344cfef 100644 --- a/src/genelab/bridges/keyboard.py +++ b/src/genelab/bridges/keyboard.py @@ -187,6 +187,12 @@ def _broadcast_hud(self, viewer: object) -> None: if not self.cfg.show_hud: return msg = f"teleop vx={self._vx:+.2f} vy={self._vy:+.2f} wz={self._wz:+.2f}" + # Genesis 1.2's outer Viewer forwards register_keybinds but not set_message_text — + # the HUD line lives on the inner pyrender viewer (same wrapper gap as the + # viewer.plugins / _viewer_plugins fallback in scene). set_text = getattr(viewer, "set_message_text", None) + if set_text is None: + inner = getattr(viewer, "_pyrender_viewer", None) + set_text = getattr(inner, "set_message_text", None) if set_text is not None: set_text(msg) diff --git a/src/genelab/entity/__init__.py b/src/genelab/entity/__init__.py index 460c688d..4cdeeb8f 100644 --- a/src/genelab/entity/__init__.py +++ b/src/genelab/entity/__init__.py @@ -3,6 +3,7 @@ from genelab.entity.articulation import Articulation, ArticulationCfg, RobotState from genelab.entity.avatar import Avatar, AvatarCfg from genelab.entity.rigid_object import RigidObject, RigidObjectCfg +from genelab.entity.root_velocity import write_root_velocity __all__ = [ "Articulation", @@ -12,4 +13,5 @@ "RigidObject", "RigidObjectCfg", "RobotState", + "write_root_velocity", ] diff --git a/src/genelab/entity/_articulation_writer.py b/src/genelab/entity/_articulation_writer.py index 70101dab..32f1a385 100644 --- a/src/genelab/entity/_articulation_writer.py +++ b/src/genelab/entity/_articulation_writer.py @@ -14,6 +14,7 @@ from genelab.actuator import ActuatorBase from genelab.entity._torch import to_tensor +from genelab.entity.root_velocity import write_root_velocity class ArticulationWriter: @@ -77,13 +78,12 @@ def write_root_state( for fn_name, value in ( ("set_pos", root_pos), ("set_quat", root_quat), - ("set_vel", root_lin_vel_w), - ("set_ang", root_ang_vel_w), ): fn = getattr(robot, fn_name, None) if fn is None: continue fn(value, envs_idx=env_ids) + write_root_velocity(robot, root_lin_vel_w, root_ang_vel_w, env_ids) def reset(self, env_ids: torch.Tensor) -> None: """Reset actuated joints to default pose + zero velocity for ``env_ids``.""" diff --git a/src/genelab/entity/root_velocity.py b/src/genelab/entity/root_velocity.py new file mode 100644 index 00000000..680b63d1 --- /dev/null +++ b/src/genelab/entity/root_velocity.py @@ -0,0 +1,69 @@ +"""Root-velocity write-back for Genesis entity handles. + +Genesis's ``RigidEntity`` API is asymmetric: ``get_vel`` / ``get_ang`` exist, but the +matching setters ``set_vel`` / ``set_ang`` exist only on FEM / tool entities. Writing a +floating base's velocity therefore goes through ``set_dofs_velocity`` on the base free +joint's 6 DoFs — indices 0–2 carry world-frame linear velocity, 3–5 world-frame angular +velocity. Every GeneLab site that overwrites root velocity (event terms, the articulation +writer, example resets) must route through :func:`write_root_velocity`; calling +``getattr(handle, "set_vel", ...)`` directly silently no-ops on rigid entities (#242). +""" + +from typing import Any + +import torch + + +def base_dof_indices(handle: Any) -> list[int] | None: + """The 6 free-joint DoF indices of ``handle``'s floating base, or ``None`` if fixed-based. + + Scans ``handle.joints`` for the first joint with ≥ 6 DoFs (same idiom as the + articulation binder's free-joint detection) and returns its first six + ``dofs_idx_local`` entries, falling back to ``dof_start`` arithmetic for handles + that don't expose the index list. + """ + joints = getattr(handle, "joints", None) or [] + for joint in joints: + if int(getattr(joint, "n_dofs", 1)) < 6: + continue + idx = getattr(joint, "dofs_idx_local", None) + if idx is not None: + idx = [int(i) for i in idx] + if len(idx) >= 6: + return idx[:6] + start = int(getattr(joint, "dof_start", 0)) + return list(range(start, start + 6)) + return None + + +def write_root_velocity( + handle: Any, + lin_vel_w: torch.Tensor, + ang_vel_w: torch.Tensor, + env_ids: torch.Tensor | None = None, +) -> bool: + """Overwrite ``handle``'s world-frame root velocity; returns ``True`` if written. + + ``lin_vel_w`` / ``ang_vel_w`` are ``(n, 3)`` aligned with ``env_ids`` (``None`` → + all envs). Rigid entities take the free-joint ``set_dofs_velocity`` path; entities + exposing direct ``set_vel`` / ``set_ang`` setters (FEM / tool entities, test fakes) + take those. ``False`` means the handle has neither — e.g. a fixed-base articulation, + which has no root velocity to write. + """ + set_dofs_velocity = getattr(handle, "set_dofs_velocity", None) + base_idx = base_dof_indices(handle) + if set_dofs_velocity is not None and base_idx is not None: + velocity = torch.cat([lin_vel_w, ang_vel_w], dim=-1) + set_dofs_velocity(velocity, base_idx, envs_idx=env_ids) + return True + wrote = False + for name, value in (("set_vel", lin_vel_w), ("set_ang", ang_vel_w)): + fn = getattr(handle, name, None) + if fn is None: + continue + try: + fn(value, envs_idx=env_ids) + except TypeError: + fn(value) + wrote = True + return wrote diff --git a/src/genelab/mdp/events.py b/src/genelab/mdp/events.py index 82b7b3b7..896fdba3 100644 --- a/src/genelab/mdp/events.py +++ b/src/genelab/mdp/events.py @@ -4,6 +4,7 @@ import torch +from genelab.entity.root_velocity import write_root_velocity from genelab.mdp._helpers import asset_articulation, asset_handle from genelab.utils.math import quat_from_euler_xyz, quat_mul @@ -72,18 +73,7 @@ def reset_root_state_uniform( if axis in velocity_range: lo, hi = velocity_range[axis] ang[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi) - set_vel = getattr(handle, "set_vel", None) - set_ang = getattr(handle, "set_ang", None) - if set_vel is not None: - try: - set_vel(vel, envs_idx=env_ids) - except TypeError: - set_vel(vel) - if set_ang is not None: - try: - set_ang(ang, envs_idx=env_ids) - except TypeError: - set_ang(ang) + write_root_velocity(handle, vel, ang, env_ids) def reset_joints_to_default( @@ -170,18 +160,7 @@ def push_by_setting_velocity( if axis in velocity_range: lo, hi = velocity_range[axis] ang[:, idx] = torch.empty(n, device=env.device).uniform_(lo, hi) - set_vel = getattr(handle, "set_vel", None) - set_ang = getattr(handle, "set_ang", None) - if set_vel is not None: - try: - set_vel(vel, envs_idx=env_ids) - except TypeError: - set_vel(vel) - if set_ang is not None: - try: - set_ang(ang, envs_idx=env_ids) - except TypeError: - set_ang(ang) + write_root_velocity(handle, vel, ang, env_ids) def randomize_terrain_params( diff --git a/tests/test_bridges.py b/tests/test_bridges.py index 60163111..081f38cf 100644 --- a/tests/test_bridges.py +++ b/tests/test_bridges.py @@ -195,3 +195,42 @@ class _HeadlessEnv: bridge.on_build(_HeadlessEnv()) # type: ignore[arg-type] bridge.pre_step(_HeadlessEnv()) # type: ignore[arg-type] + + +def test_keyboard_hud_falls_back_to_inner_pyrender_viewer() -> None: + """Genesis 1.2's outer Viewer forwards register_keybinds but not set_message_text, + so the HUD write must reach the inner ``_pyrender_viewer`` (same wrapper gap as + the ``viewer.plugins`` / ``_viewer_plugins`` fallback in scene).""" + bridge = KeyboardTwistBridge(KeyboardTwistBridgeCfg()) + + class _InnerViewer: + def __init__(self) -> None: + self.messages: list[str] = [] + + def set_message_text(self, text: str) -> None: + self.messages.append(text) + + class _OuterViewer: + def __init__(self) -> None: + self._pyrender_viewer = _InnerViewer() + + viewer = _OuterViewer() + bridge._broadcast_hud(viewer) + assert viewer._pyrender_viewer.messages + assert viewer._pyrender_viewer.messages[0].startswith("teleop") + + +def test_keyboard_hud_prefers_direct_setter() -> None: + """A viewer exposing set_message_text directly (older Genesis, fakes) is used as-is.""" + bridge = KeyboardTwistBridge(KeyboardTwistBridgeCfg()) + + class _DirectViewer: + def __init__(self) -> None: + self.messages: list[str] = [] + + def set_message_text(self, text: str) -> None: + self.messages.append(text) + + viewer = _DirectViewer() + bridge._broadcast_hud(viewer) + assert viewer.messages and viewer.messages[0].startswith("teleop") diff --git a/tests/test_genesis_api_contract.py b/tests/test_genesis_api_contract.py new file mode 100644 index 00000000..13aee444 --- /dev/null +++ b/tests/test_genesis_api_contract.py @@ -0,0 +1,87 @@ +"""Contract pins for every Genesis method GeneLab calls through a ``getattr`` guard. + +GeneLab's write seams duck-type against the Genesis handle +(``getattr(handle, "set_...", None)``) so unit-test fakes and multiple entity types +stay supported. The cost of that pattern is #242: a method the upstream API never had +(or renamed) silently no-ops instead of raising — ``push_by_setting_velocity`` was dead +for every rigid articulation because ``RigidEntity`` never had ``set_vel`` / ``set_ang``. + +These tests pin every ``getattr``-guarded name against the *installed* Genesis, so an +upstream rename surfaces as a test failure at the next version bump instead of a silent +physics change. Import-only (no ``gs.init`` scene build), so they also run on headless CI. +""" + +import pytest + +pytest.importorskip("genesis") + +# Every RigidEntity method reached via ``getattr(handle, name, None)`` in src/genelab +# and examples/ (grep for ``getattr(`` + ``"set_/get_/control_"`` to regenerate). +# Names in dead *fallback* arms of ``x or y`` chains are deliberately not pinned +# (``set_dofs_friction``, ``set_dofs_position_target``) — their primary arm is. +RIGID_ENTITY_SEAMS = ( + # root / DoF write seams + "set_pos", + "set_quat", + "set_dofs_position", + "set_dofs_velocity", + "set_dofs_kp", + "set_dofs_kv", + "set_dofs_force_range", + "set_dofs_armature", + "set_dofs_frictionloss", + "set_friction_ratio", + "set_mass_shift", + "control_dofs_position", + "control_dofs_force", + "control_dofs_velocity", + # read seams + "get_dofs_force", + "get_dofs_control_force", + "get_dofs_armature", + "get_dofs_limit", + "get_links_inertial_mass", + "get_links_net_contact_force", + "get_contacts", + "get_vel", + "get_ang", +) + + +def test_rigid_entity_seams_exist() -> None: + from genesis.engine.entities import RigidEntity + + missing = [name for name in RIGID_ENTITY_SEAMS if not hasattr(RigidEntity, name)] + assert not missing, ( + f"Genesis RigidEntity no longer has {missing} — the GeneLab call sites guarding " + f"these with getattr(..., None) would silently no-op (see #242). Update the seam " + f"or this pin." + ) + + +def test_free_joint_surface_for_root_velocity_writes() -> None: + """``write_root_velocity`` needs joints with n_dofs / dofs_idx_local for free-joint discovery.""" + from genesis.engine.entities.rigid_entity.rigid_joint import RigidJoint + + for attr in ("n_dofs", "dofs_idx_local", "dof_start"): + assert hasattr(RigidJoint, attr), f"RigidJoint lost {attr}" + + +def test_viewer_hud_seam() -> None: + """The teleop HUD writes through the inner pyrender viewer's set_message_text.""" + viewer_mod = pytest.importorskip("genesis.ext.pyrender.viewer") + assert hasattr(viewer_mod.Viewer, "set_message_text") + assert hasattr(viewer_mod.Viewer, "register_keybinds") + + +def test_rigid_solver_gravity_seam() -> None: + """``mdp.dr.gravity`` writes per-env gravity through ``sim.rigid_solver.set_gravity``.""" + import genesis as gs + + try: + if not getattr(gs, "_initialized", False): + gs.init(backend=gs.cpu, logging_level="error") + from genesis.engine.solvers import RigidSolver + except Exception as exc: # noqa: BLE001 - headless CI may fail at init, not import + pytest.skip(f"Genesis solver import needs a runtime: {exc}") + assert hasattr(RigidSolver, "set_gravity") diff --git a/tests/test_root_velocity.py b/tests/test_root_velocity.py new file mode 100644 index 00000000..ceb7a6f8 --- /dev/null +++ b/tests/test_root_velocity.py @@ -0,0 +1,227 @@ +"""Tests for ``genelab.entity.root_velocity`` — the #242 regression seam. + +Genesis's ``RigidEntity`` exposes ``get_vel`` / ``get_ang`` but no ``set_vel`` / +``set_ang`` (those live on FEM / tool entities only), so every root-velocity write +must route through ``set_dofs_velocity`` on the base free joint's 6 DoFs. The fakes +here mirror that Genesis 1.2 API surface — deliberately *without* ``set_vel`` / +``set_ang`` — so a regression back to the old ``getattr`` idiom fails these tests +instead of silently no-opping. +""" + +from typing import Any + +import torch + +from genelab.entity._articulation_writer import ArticulationWriter +from genelab.entity.root_velocity import base_dof_indices, write_root_velocity +from genelab.mdp.events import push_by_setting_velocity, reset_root_state_uniform + + +class _FakeJoint: + def __init__(self, n_dofs: int, dofs_idx_local: list[int]) -> None: + self.n_dofs = n_dofs + self.dofs_idx_local = dofs_idx_local + + +class _FakeRigidEntity: + """Genesis-1.2-shaped rigid handle: free joint + hinges, no ``set_vel`` / ``set_ang``.""" + + def __init__(self, free_base: bool = True) -> None: + hinge_offset = 6 if free_base else 0 + self.joints = ([_FakeJoint(6, [0, 1, 2, 3, 4, 5])] if free_base else []) + [ + _FakeJoint(1, [hinge_offset]), + _FakeJoint(1, [hinge_offset + 1]), + ] + self.dofs_velocity_calls: list[tuple[torch.Tensor, list[int], Any]] = [] + self.pos_calls: list[torch.Tensor] = [] + self.quat_calls: list[torch.Tensor] = [] + + def set_dofs_velocity( + self, velocity: torch.Tensor, dofs_idx_local: Any = None, envs_idx: Any = None + ) -> None: + self.dofs_velocity_calls.append((velocity.clone(), list(dofs_idx_local), envs_idx)) + + def set_pos(self, value: torch.Tensor, envs_idx: Any = None) -> None: + self.pos_calls.append(value.clone()) + + def set_quat(self, value: torch.Tensor, envs_idx: Any = None) -> None: + self.quat_calls.append(value.clone()) + + +class _FakeSoftEntity: + """FEM / tool-entity shape: direct ``set_vel`` / ``set_ang``, no DoF API.""" + + def __init__(self) -> None: + self.vel: torch.Tensor | None = None + self.ang: torch.Tensor | None = None + + def set_vel(self, value: torch.Tensor, envs_idx: Any = None) -> None: + self.vel = value.clone() + + def set_ang(self, value: torch.Tensor, envs_idx: Any = None) -> None: + self.ang = value.clone() + + +# ------------------------------------------------------------- base_dof_indices + + +def test_base_dof_indices_free_base() -> None: + assert base_dof_indices(_FakeRigidEntity()) == [0, 1, 2, 3, 4, 5] + + +def test_base_dof_indices_fixed_base_is_none() -> None: + assert base_dof_indices(_FakeRigidEntity(free_base=False)) is None + + +def test_base_dof_indices_no_joints_is_none() -> None: + assert base_dof_indices(object()) is None + + +# ------------------------------------------------------------ write_root_velocity + + +def test_write_root_velocity_routes_through_free_joint_dofs() -> None: + handle = _FakeRigidEntity() + lin = torch.tensor([[1.0, 2.0, 3.0]]) + ang = torch.tensor([[0.1, 0.2, 0.3]]) + env_ids = torch.tensor([4]) + assert write_root_velocity(handle, lin, ang, env_ids) is True + (velocity, idx, envs_idx) = handle.dofs_velocity_calls[0] + assert torch.equal(velocity, torch.tensor([[1.0, 2.0, 3.0, 0.1, 0.2, 0.3]])) + assert idx == [0, 1, 2, 3, 4, 5] + assert torch.equal(envs_idx, env_ids) + + +def test_write_root_velocity_fixed_base_writes_nothing() -> None: + handle = _FakeRigidEntity(free_base=False) + zeros = torch.zeros(1, 3) + assert write_root_velocity(handle, zeros, zeros, torch.tensor([0])) is False + assert handle.dofs_velocity_calls == [] + + +def test_write_root_velocity_direct_setter_fallback() -> None: + handle = _FakeSoftEntity() + lin = torch.tensor([[1.0, 0.0, 0.0]]) + ang = torch.tensor([[0.0, 0.0, 2.0]]) + assert write_root_velocity(handle, lin, ang, torch.tensor([0])) is True + assert handle.vel is not None and torch.equal(handle.vel, lin) + assert handle.ang is not None and torch.equal(handle.ang, ang) + + +# ------------------------------------------------------------------ event terms + + +class _FakeEntityCfg: + init_pos = (0.0, 0.0, 0.5) + init_quat = (1.0, 0.0, 0.0, 0.0) + + +class _FakeArticulation: + cfg = _FakeEntityCfg() + + +class _FakeEnv: + def __init__(self, robot: Any) -> None: + self.robot = robot + self.articulation = _FakeArticulation() + self.device = "cpu" + + +def test_push_by_setting_velocity_writes_sampled_base_velocity() -> None: + """#242: the push must land in the free-joint DoFs of a handle without set_vel/set_ang.""" + torch.manual_seed(0) + robot = _FakeRigidEntity() + env = _FakeEnv(robot) + env_ids = torch.arange(32) + push_by_setting_velocity(env, env_ids, velocity_range={"x": (0.5, 1.0), "yaw": (-0.2, -0.1)}) + assert len(robot.dofs_velocity_calls) == 1 + velocity, idx, envs_idx = robot.dofs_velocity_calls[0] + assert velocity.shape == (32, 6) + assert idx == [0, 1, 2, 3, 4, 5] + assert torch.equal(envs_idx, env_ids) + assert torch.all(velocity[:, 0] >= 0.5) and torch.all(velocity[:, 0] <= 1.0) + assert torch.all(velocity[:, 5] >= -0.2) and torch.all(velocity[:, 5] <= -0.1) + # Unspecified axes overwrite to zero (Isaac Lab parity). + assert torch.equal(velocity[:, 1:5], torch.zeros(32, 4)) + + +def test_push_by_setting_velocity_empty_env_ids_is_noop() -> None: + robot = _FakeRigidEntity() + push_by_setting_velocity(_FakeEnv(robot), torch.tensor([], dtype=torch.long)) + assert robot.dofs_velocity_calls == [] + + +def test_reset_root_state_uniform_writes_base_velocity() -> None: + torch.manual_seed(0) + robot = _FakeRigidEntity() + env = _FakeEnv(robot) + env_ids = torch.arange(16) + reset_root_state_uniform(env, env_ids, velocity_range={"z": (0.1, 0.2)}) + assert len(robot.pos_calls) == 1 # pose write still happens + assert len(robot.dofs_velocity_calls) == 1 + velocity, idx, _ = robot.dofs_velocity_calls[0] + assert idx == [0, 1, 2, 3, 4, 5] + assert torch.all(velocity[:, 2] >= 0.1) and torch.all(velocity[:, 2] <= 0.2) + + +# ------------------------------------------------------------------ writer seam + + +def test_writer_write_root_state_routes_velocity_via_free_joint() -> None: + robot = _FakeRigidEntity() + writer = ArticulationWriter( + robot, + actuated_dof_idx=torch.tensor([6, 7]), + default_joint_pos=torch.zeros(2), + joint_pos_target=torch.zeros(1, 2), + actuators={}, + device="cpu", + ) + env_ids = torch.tensor([0]) + writer.write_root_state( + torch.zeros(1, 3), + torch.tensor([[1.0, 0.0, 0.0, 0.0]]), + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[0.1, 0.2, 0.3]]), + env_ids, + ) + assert len(robot.pos_calls) == 1 and len(robot.quat_calls) == 1 + velocity, idx, envs_idx = robot.dofs_velocity_calls[0] + assert torch.equal(velocity, torch.tensor([[1.0, 2.0, 3.0, 0.1, 0.2, 0.3]])) + assert idx == [0, 1, 2, 3, 4, 5] + assert torch.equal(envs_idx, env_ids) + + +# ------------------------------------------------------------- real Genesis seam + + +def test_write_root_velocity_round_trips_on_real_genesis(genesis_runtime: Any) -> None: + """The definitive #242 check: write through the free joint, read back via get_vel/get_ang.""" + del genesis_runtime # fixture only guards EGL/runtime availability + from genelab.configs import InteractiveSceneCfg, SimulationCfg + from genelab.entity import RigidObjectCfg + from genelab.scene import InteractiveScene + + sim_cfg = SimulationCfg(num_envs=2, dt=0.01, substeps=1, vis=False, gpu=False) + scene_cfg = InteractiveSceneCfg( + env_spacing=(2.0, 2.0), + # ``fixed=False``: only a floating box has the free joint this write targets. + entities={ + "box": RigidObjectCfg( + morph="box", size=(0.1, 0.1, 0.1), init_pos=(0, 0, 1.0), fixed=False + ) + }, + ) + scene = InteractiveScene(sim_cfg, scene_cfg, device_hint="cpu") + try: + scene.build() + handle = scene.rigid_objects["box"].gs_handle + lin = torch.tensor([[1.0, 2.0, 3.0], [0.5, 0.0, 0.0]]) + ang = torch.tensor([[0.1, 0.2, 0.3], [0.0, 0.0, 0.5]]) + assert write_root_velocity(handle, lin, ang, torch.tensor([0, 1])) is True + got_lin = torch.as_tensor(handle.get_vel()).cpu() + got_ang = torch.as_tensor(handle.get_ang()).cpu() + assert torch.allclose(got_lin, lin, atol=1e-5) + assert torch.allclose(got_ang, ang, atol=1e-5) + finally: + scene.close()