Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion robolab/core/environments/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,8 @@ def generate_task_env_cfg(task_class: Task,
gripper_closure_cfg: dict | None = None,
lazy_sensor_update: bool = True,
ee_recorder_bodies: dict[str, str] | None = None,
object_state_obs: bool = False) -> Type[RobolabDefaultEnvCfg]:
object_state_obs: bool = False,
solver_iterations: tuple[int, int] | None = None) -> Type[RobolabDefaultEnvCfg]:
"""
Generate a complete task environment configuration class.

Expand All @@ -152,6 +153,8 @@ def generate_task_env_cfg(task_class: Task,
meters), ``<object>_quat`` (world-frame w, x, y, z), and
``<object>_vel`` (world-frame) terms for every entry of the
task's ``contact_object_list`` (minus fixtures). Default False.
solver_iterations: Optional scene solver limits (position, velocity),
applied after global defaults so robot iteration requests are not capped.

Returns:
A complete environment configuration class
Expand Down Expand Up @@ -194,6 +197,12 @@ class GeneratedTaskEnvCfg(RobolabDefaultEnvCfg):
def __post_init__(self):
super().__post_init__() # Set all defaults first

if solver_iterations is not None:
for axis, count in zip(("position", "velocity"), solver_iterations):
for field in (f"num_{axis}_iterations", f"max_{axis}_iteration_count"):
if hasattr(self.sim.physx, field):
setattr(self.sim.physx, field, count)

self.episode_length_s: int = task_class.episode_length_s
self.decimation: int = decimation
self.sim.dt: int = dt
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ def is_robot_attached(camera_cls):
dt=dt,
render_interval=render_interval,
decimation=decimation,
solver_iterations=(128, 4),
seed=1,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ def is_robot_attached(camera_cls):
dt=dt,
render_interval=render_interval,
decimation=decimation,
solver_iterations=(128, 4),
seed=1,
)

Expand Down
40 changes: 40 additions & 0 deletions tests/test_galbot_solver_iterations.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Solver limits after task registration and runtime configuration parsing."""

import pytest

from robolab.core.environments.config import parse_env_cfg
from robolab.registrations.droid.auto_env_registrations_jointpos import auto_register_droid_envs
from robolab.registrations.galbot.auto_env_registrations_abs_ik import auto_register_galbot_abs_ik_envs
from robolab.registrations.galbot.auto_env_registrations_jointpos import auto_register_galbot_envs


@pytest.mark.parametrize(
"register, kwargs, expected",
[
(auto_register_galbot_envs, {"action": "whole_body"}, (128, 4)),
(auto_register_galbot_envs, {"action": "arms"}, (128, 4)),
(auto_register_galbot_abs_ik_envs, {}, (128, 4)),
(auto_register_droid_envs, {}, (32, 1)),
],
)
def test_registered_task_solver_limits(register, kwargs, expected):
postfix = f"SolverLimits{register.__name__}{kwargs.get('action', '')}"
if register is auto_register_droid_envs:
register(task="BananaInBowlTask")
postfix = ""
else:
register(task="BananaInBowlTask", env_postfix=postfix, **kwargs)
cfg = parse_env_cfg(f"BananaInBowlTask{postfix}", num_envs=1)
if expected == (128, 4):
articulation = cfg.scene.robot.spawn.articulation_props
assert articulation.solver_position_iteration_count == 128
assert articulation.solver_velocity_iteration_count == 4
for axis, count in zip(("position", "velocity"), expected):
physx = cfg.sim.physx
field = f"max_{axis}_iteration_count"
if not hasattr(physx, field):
field = f"num_{axis}_iterations"
assert getattr(physx, field) == count