From d0e9066075ccf2376602ee7cbe5100b73831b5b6 Mon Sep 17 00:00:00 2001
From: ruziniuuuuu <ruziniuuuuu@gmail.com>
Date: Fri, 11 Sep 2026 23:19:13 +0800
Subject: [PATCH] fix: preserve Galbot solver iterations in task environments

Signed-off-by: ruziniuuuuu <ruziniuuuuu@gmail.com>
---
 robolab/core/environments/config.py           | 11 ++++-
 .../galbot/auto_env_registrations_abs_ik.py   |  1 +
 .../galbot/auto_env_registrations_jointpos.py |  1 +
 tests/test_galbot_solver_iterations.py        | 40 +++++++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_galbot_solver_iterations.py

diff --git a/robolab/core/environments/config.py b/robolab/core/environments/config.py
index d62e2c75..f712bb0a 100644
--- a/robolab/core/environments/config.py
+++ b/robolab/core/environments/config.py
@@ -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.
 
@@ -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
@@ -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
diff --git a/robolab/registrations/galbot/auto_env_registrations_abs_ik.py b/robolab/registrations/galbot/auto_env_registrations_abs_ik.py
index 0dfc1cdd..3292d2ee 100644
--- a/robolab/registrations/galbot/auto_env_registrations_abs_ik.py
+++ b/robolab/registrations/galbot/auto_env_registrations_abs_ik.py
@@ -92,6 +92,7 @@ def is_robot_attached(camera_cls):
         dt=dt,
         render_interval=render_interval,
         decimation=decimation,
+        solver_iterations=(128, 4),
         seed=1,
     )
 
diff --git a/robolab/registrations/galbot/auto_env_registrations_jointpos.py b/robolab/registrations/galbot/auto_env_registrations_jointpos.py
index d6ed1326..eafc9cf1 100644
--- a/robolab/registrations/galbot/auto_env_registrations_jointpos.py
+++ b/robolab/registrations/galbot/auto_env_registrations_jointpos.py
@@ -101,6 +101,7 @@ def is_robot_attached(camera_cls):
         dt=dt,
         render_interval=render_interval,
         decimation=decimation,
+        solver_iterations=(128, 4),
         seed=1,
     )
 
diff --git a/tests/test_galbot_solver_iterations.py b/tests/test_galbot_solver_iterations.py
new file mode 100644
index 00000000..829d56d7
--- /dev/null
+++ b/tests/test_galbot_solver_iterations.py
@@ -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