Physics working in slow motion for non-agents when using multiple envs

[elKuston]

Hello! I am designing a DirectMARLEnv where the agents are drones, and there are UGVs as well, which are not agents, but rather "dynamic parts of the enviroment". I programmatically calculate the movement of the UGVs and apply it in the _pre_physics_step() as following, which is essentially a simple implementation of Pure Pursuit:

def _pre_physics_step(self, actions: dict[str, torch.Tensor]) -> None:
        self.actions = actions
        # (N_envs, n_drones, joints)
        self.stacked_actions = torch.stack(
            [self.actions[agent_name] for agent_name in self.cfg.possible_agents], dim=1
        )
        self.visualize_markers()
        self.control_ugv()



def compute_jackal_commands(self):
        """
        Calcula comandos v y omega para N entornos simultáneamente.
        """
        # 0. Si ha llegado a las coordenadas, calcular nuevas
        dist_to_goal = torch.norm(
            self.ugv_goal_pos[:, :2] - self.ugv.data.root_pos_w[:, :2],
            dim=-1,
        )

        reached_indices = (dist_to_goal < 0.5).nonzero(as_tuple=True)[0]

        if len(reached_indices) > 0:
            num_new = len(reached_indices)
            new_goals = torch.rand((num_new, 3), device=self.device) * 50.0 - 25.0

            # Forzar la Z a 0 (el Jackal no vuela... aún)
            new_goals[:, 2] = 0.0

            # Asignar solo a los entornos que lo necesitan
            self.ugv_goal_pos[reached_indices] = (
                new_goals + self.scene.env_origins[reached_indices]
            )

        # 1. Vector de error en coordenadas globales
        error_world = self.ugv_goal_pos - self.ugv.data.root_pos_w

        # 2. Transformar el error al sistema LOCAL del Jackal
        # (x_local apunta hacia adelante, y_local hacia la izquierda)
        error_local = math_utils.quat_apply_inverse(
            self.ugv.data.root_link_quat_w, error_world
        )

        # 3. Parámetros de control
        v_ref = 2.0  # m/s
        k_yaw = 4.5  # Ganancia de giro
        dist_threshold = 0.5  # Distancia para considerar waypoint alcanzado

        # 4. Cálculo de comandos
        # El ángulo hacia el objetivo en el marco local
        heading_error = torch.atan2(error_local[:, 1], error_local[:, 0])

        # Velocidad lineal: se reduce si el giro es muy pronunciado
        v_cmd = v_ref * torch.cos(heading_error).clamp(min=0.0)

        # Velocidad angular: proporcional al error de heading
        omega_cmd = k_yaw * heading_error

        # Detenerse si estamos cerca del objetivo
        dist_to_target = torch.norm(error_local[:, :2], dim=-1)
        v_cmd = torch.where(dist_to_target < dist_threshold, 0.0, v_cmd)
        omega_cmd = torch.where(dist_to_target < dist_threshold, 0.0, omega_cmd)

        return v_cmd, omega_cmd

    def control_ugv(self):
        # Parámetros físicos del Jackal escalado (aproximados)
        wheel_base = 0.6  # Distancia lateral entre ruedas (m)
        wheel_radius = 0.3  # Radio de la rueda escalada (m)

        # v = (r * (w_r + w_l)) / 2  --> Velocidad lineal
        # w = (r * (w_r - w_l)) / L  --> Velocidad angular

        v, omega = self.compute_jackal_commands()

        v_left = (v - omega * wheel_base / 2.0) / wheel_radius
        v_right = (v + omega * wheel_base / 2.0) / wheel_radius

        # El Jackal suele tener este orden de juntas: FL, FR, RL, RR
        # Replicamos las velocidades para las ruedas traseras
        joint_vel_targets = torch.stack([v_left, v_right, v_left, v_right], dim=-1)

        # Enviar a la GPU
        self.ugv.set_joint_velocity_target(joint_vel_targets)

While defining the UGVs in the _setup_scene() as:

def _setup_scene(self):
        self._setup_common()
        self.ugv = Articulation(self.cfg.ugv_cfg)
        self.ugv_goal_pos = torch.zeros((self.num_envs, 3), device=self.device)
        self.scene.articulations["ugv"] = self.ugv

With its config being:

UGV_CONFIG = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        usd_path="[...]/jackal.usd",
    ),
    actuators={
        "wheel_acts": ImplicitActuatorCfg(
            joint_names_expr=[".*"], damping=1e5, stiffness=0
        )
    },
)

For the UGVs' USD I have used both the nvidia-provided Clearpath Jackal USD, as well as a custom, very simple robot (similar to what one could find in the Building a simple robot Nvidia tutorial)

The issue

Drones (that is, agents), work fine in every condition. UGVs, on the other hand, work as expected when using a low number of environments (up to 8), but when moving to a larger number of envs, physics start to work in "slow motion". When instantiating the UGVs above the ground level, for example, one can see them fall in slow motion. The joints (wheels), however, move at the appropiate speed, although friction doesn't allow the vehicle to move. Drones, however, work fine, both in the speed of joints, and when it comes to physics.

I've already tried modifying many gpu parameters, such as gpu_max_rigid_contact_count, gpu_max_rigid_patch_count, or gpu_heap_capacity, but I didn't manage to solve the issue.

I've been stuck with this for almost two weeks. Is there something I am missing? What could be the cause for this and how could I fix this?
Thank you very much for your support

PD: My build is composed of an RTX3080, 32GB of RAM and an i5 14600KF