Skip to content

Comments

Switching the wrench compose to rely on two buffers (local and global)#1

Open
AntoineRichard wants to merge 11 commits intogbionics:fix/v2.3.2/wrench_composerfrom
AntoineRichard:antoiner/wrench_composer_dual
Open

Switching the wrench compose to rely on two buffers (local and global)#1
AntoineRichard wants to merge 11 commits intogbionics:fix/v2.3.2/wrench_composerfrom
AntoineRichard:antoiner/wrench_composer_dual

Conversation

@AntoineRichard
Copy link

Description

Following up on the comment i made on your PR, I think we need two buffers to handle permanent local and global wrenches properly.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Checklist

  • I have read and understood the contribution guidelines
  • I have run the pre-commit checks with ./isaaclab.sh --format
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the changelog and the corresponding version in the extension's config/extension.toml file
  • I have added my name to the CONTRIBUTORS.md or my name already exists there

@AntoineRichard AntoineRichard changed the title committing dual buffer. Switching the wrench compose to rely on two buffers (local and global) Feb 16, 2026
…nches would not be updated properly when the object would move.
@AntoineRichard
Copy link
Author

@LoreMoretti could you take a look and let me know if this version of the PR works for you guys. I extended the test coverage quite drastically, but if you can think of some other tests let me know. This code is critical and should be bullet proof!

Copy link

@LoreMoretti LoreMoretti left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! I'll run a training to check whether our use case is working fine now.

@LoreMoretti
Copy link

I'll run a training to check whether our use case is working fine now.

I ran a training. I get a lot of the following warnings:

Transform has scale 0.001000, world position (-9798498123776.000000, -75850464624640.000000, 83705188057088.000000), with bounding box min (-53.195415, -46.511654, -112.450203) and max (55.791599, 39.303677, 16.042425)
2026-02-17T17:04:03Z [80,729ms] [Warning] [rtx.scenedb.plugin] Instance 114389 of geometry "/__Prototype_1992889926339252044/<link_name>/node_STL_BINARY_/mesh" has bounding box dimensions after transformation that exceed the recommended extents limit of 1099511627776.000000.

Which do not appear, if I use the v2.3.2 tag of IsaacLab. For a lot of <link_names>.

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 18, 2026

Hum something is exploding for sure. I could be that the torque correction is problematic when objects are far from the origin.

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 18, 2026

@LoreMoretti I ran some more test trying the wrench correction at 2000 meters from the origin and comparing to 1 meter and the results are inline. Would you mind sharing a bit more about what's creating this issue? Something like the position at which the force is applied, the frame in which it's applied, and the magnitude?

Would it be possible that since the wrenches are now working as expected you're getting increasingly larger wrenches as the robot moves away from the application point of the force. The global force doesn't track the body pose. So if you apply a global force on a body it will default to applying the force in (0,0,0). Which could apply very large torques to the robots far from the application point.

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 18, 2026

We could make the global forces track the robot position but we'd need another set of buffers to enable that. The feature could be: if no position is provided, global forces are applied at the CoM position of the bodies, forces do not generate a resulting torque, but their magnitude/orientation is expressed in the global frame.

@LoreMoretti
Copy link

LoreMoretti commented Feb 18, 2026

@LoreMoretti I ran some more test trying the wrench correction at 2000 meters from the origin and comparing to 1 meter and the results are inline. Would you mind sharing a bit more about what's creating this issue? Something like the position at which the force is applied, the frame in which it's applied, and the magnitude?

Would it be possible that since the wrenches are now working as expected you're getting increasingly larger wrenches as the robot moves away from the application point of the force. The global force doesn't track the body pose. So if you apply a global force on a body it will default to applying the force in (0,0,0). Which could apply very large torques to the robots far from the application point.

I have to investigate more on the issue. What I can say is that it was working with no warnings also in v2.3.1, where there was not the wrench composer. I did not investigate on how forces were handled before the wrench composer. Weren't global forces tracking body position there?

@AntoineRichard
Copy link
Author

They were I just pushed a fix for that:

    """Set external force and torque to apply on the asset's bodies in their local frame.

    For many applications, we want to keep the applied external force on rigid bodies constant over a period of
    time (for instance, during the policy control). This function allows us to store the external force and torque
    into buffers which are then applied to the simulation at every step. Optionally, set the position to apply the
    external wrench at (in the local link frame of the bodies).

    .. caution::
        This method clears **all** internal wrench buffers before writing the provided values.
        Any previously set forces or torques (local or global) are discarded. To accumulate
        on top of existing values, use ``permanent_wrench_composer.add_forces_and_torques`` instead.

    .. note::
        This function does not apply the external wrench to the simulation. It only fills the buffers with
        the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function
        right before the simulation step.

    Args:
        forces: External forces. Shape is (len(env_ids), len(body_ids), 3).
            When ``is_global=False``, forces are in the bodies' local frame.
            When ``is_global=True``, forces are in the world frame.
        torques: External torques. Shape is (len(env_ids), len(body_ids), 3).
            When ``is_global=False``, torques are in the bodies' local frame.
            When ``is_global=True``, torques are in the world frame.
        positions: Application points for forces. Shape is (len(env_ids), len(body_ids), 3).
            Defaults to None.
            When ``is_global=False``, positions are local offsets from the link frame.
            When ``is_global=True``, positions are world-frame coordinates. If None,
            forces are applied at the body's center of mass (no positional torque).
        body_ids: Body indices to apply external wrench to. Defaults to None (all bodies).
        env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
        is_global: Whether forces and torques are in the global (world) frame. Defaults to False.
    """

@AntoineRichard
Copy link
Author

Let me know if these changes do help with your training.

@LoreMoretti
Copy link

if no position is provided, global forces are applied at the CoM position of the bodies, forces do not generate a resulting torque, but their magnitude/orientation is expressed in the global frame.

Indeed this is the use case I am working on (with also 0 torques)

@LoreMoretti
Copy link

Early feedback: the behavior is different yet. I will try to provide more details as soon as I can.

@LoreMoretti
Copy link

The following is the reward of my training with 3 different version of IsaacLab:

image
  • Yellow curve is v2.3.1

  • Light blue is my branch before this PR

  • Fuchsia is this PR

Yellow and light blue curve are overlapped.

The use case is apply global forces (torques set to 0) without offsets at 2 bodies.

Your code looks good to me, I am still trying to figure out why then the behavior is different compared to previous versions of IsaacLab.

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 18, 2026

Just made these tests to compare to PhysX it looks good

It creates two views one for raw PhysX calls, one for composer calls and checks that the behavior is the same.

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Integration tests comparing WrenchComposer output vs raw PhysX apply_forces_and_torques_at_position.

Two identical rigid objects are placed in the same scene. One uses the WrenchComposer path
(set_forces_and_torques → write_data_to_sim → compose → PhysX apply with is_global=False),
the other uses the raw PhysX API directly (apply_forces_and_torques_at_position with matching
is_global flag). After N steps, both objects should have identical velocities.
"""

"""Launch Isaac Sim Simulator first."""

from isaaclab.app import AppLauncher

# launch omniverse app
simulation_app = AppLauncher(headless=True).app

"""Rest everything follows."""

import math

import pytest
import torch

import isaaclab.sim as sim_utils
from isaaclab.assets import RigidObject, RigidObjectCfg
from isaaclab.sim import build_simulation_context
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR


def generate_dual_cube_scene(
    num_cubes: int = 1,
    height: float = 1.0,
    device: str = "cuda:0",
    initial_rot: tuple[float, ...] | None = None,
) -> tuple[RigidObject, RigidObject]:
    """Generate a scene with two sets of cubes: one for the composer path, one for raw PhysX.

    Both sets share the same spawn config and initial state (except a Y offset to avoid overlap).

    Args:
        num_cubes: Number of cubes per group (environments).
        height: Spawn height.
        device: Simulation device.
        initial_rot: Initial quaternion (w, x, y, z). Defaults to identity.

    Returns:
        Tuple of (cube_composer, cube_raw) RigidObject instances.
    """
    if initial_rot is None:
        initial_rot = (1.0, 0.0, 0.0, 0.0)

    # Create Xform prims for both groups
    for i in range(num_cubes):
        origin_composer = (i * 2.0, 0.0, height)
        origin_raw = (i * 2.0, 3.0, height)  # Y offset to avoid overlap
        sim_utils.create_prim(f"/World/Composer_{i}", "Xform", translation=origin_composer)
        sim_utils.create_prim(f"/World/Raw_{i}", "Xform", translation=origin_raw)

    spawn_cfg = sim_utils.UsdFileCfg(
        usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(),
    )

    cube_composer_cfg = RigidObjectCfg(
        prim_path="/World/Composer_.*/Object",
        spawn=spawn_cfg,
        init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height), rot=initial_rot),
    )
    cube_composer = RigidObject(cfg=cube_composer_cfg)

    cube_raw_cfg = RigidObjectCfg(
        prim_path="/World/Raw_.*/Object",
        spawn=spawn_cfg,
        init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3.0, height), rot=initial_rot),
    )
    cube_raw = RigidObject(cfg=cube_raw_cfg)

    return cube_composer, cube_raw


N_STEPS = 50
FORCE_MAGNITUDE = 10.0
TORQUE_MAGNITUDE = 1.0
# 45 degrees about Z: (cos(22.5°), 0, 0, sin(22.5°))
ROT_45_Z = (math.cos(math.pi / 8), 0.0, 0.0, math.sin(math.pi / 8))


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_local_force(device):
    """Baseline: local force at identity orientation. Composer and raw PhysX should match exactly."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device)

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Composer path: local force +X
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        forces[..., 0] = FORCE_MAGNITUDE
        torques = torch.zeros(1, len(body_ids), 3, device=device)

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            body_ids=body_ids,
            is_global=False,
        )

        # Raw PhysX data (flattened for PhysX view API)
        raw_forces = torch.zeros(1, 3, device=device)
        raw_forces[:, 0] = FORCE_MAGNITUDE
        raw_torques = torch.zeros(1, 3, device=device)
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()  # no-op (composer inactive)
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=None,
                indices=raw_indices,
                is_global=False,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Compare velocities
        torch.testing.assert_close(
            cube_composer.data.root_lin_vel_w,
            cube_raw.data.root_lin_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )
        # Both should have ~zero angular velocity (force at CoM, no torque)
        torch.testing.assert_close(
            cube_composer.data.root_ang_vel_w,
            torch.zeros(1, 3, device=device),
            rtol=0.0,
            atol=1e-4,
        )
        torch.testing.assert_close(
            cube_raw.data.root_ang_vel_w,
            torch.zeros(1, 3, device=device),
            rtol=0.0,
            atol=1e-4,
        )


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_global_force(device):
    """Global force with non-identity rotation (45 deg Z). Rotation matters for frame conversion."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(
            num_cubes=1, device=device, initial_rot=ROT_45_Z
        )

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Composer path: global force +X
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        forces[..., 0] = FORCE_MAGNITUDE
        torques = torch.zeros(1, len(body_ids), 3, device=device)

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            body_ids=body_ids,
            is_global=True,
        )

        # Raw PhysX data
        raw_forces = torch.zeros(1, 3, device=device)
        raw_forces[:, 0] = FORCE_MAGNITUDE
        raw_torques = torch.zeros(1, 3, device=device)
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=None,
                indices=raw_indices,
                is_global=True,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Linear velocities should match (same global force, same mass)
        torch.testing.assert_close(
            cube_composer.data.root_lin_vel_w,
            cube_raw.data.root_lin_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_local_force_at_position(device):
    """Local force at a local offset. Both paths should produce identical cross-product torque."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device)

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Local force +X at local offset +0.5m Y
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        forces[..., 0] = FORCE_MAGNITUDE
        torques = torch.zeros(1, len(body_ids), 3, device=device)
        positions = torch.zeros(1, len(body_ids), 3, device=device)
        positions[..., 1] = 0.5  # +0.5m Y offset in local frame

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            positions=positions,
            body_ids=body_ids,
            is_global=False,
        )

        # Raw PhysX data (local force at local position)
        raw_forces = torch.zeros(1, 3, device=device)
        raw_forces[:, 0] = FORCE_MAGNITUDE
        raw_torques = torch.zeros(1, 3, device=device)
        raw_positions = torch.zeros(1, 3, device=device)
        raw_positions[:, 1] = 0.5
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=raw_positions,
                indices=raw_indices,
                is_global=False,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Both linear and angular velocities should match
        torch.testing.assert_close(
            cube_composer.data.root_lin_vel_w,
            cube_raw.data.root_lin_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )
        torch.testing.assert_close(
            cube_composer.data.root_ang_vel_w,
            cube_raw.data.root_ang_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )

        # Sanity: angular velocity should be nonzero (cross-product torque)
        assert torch.abs(cube_composer.data.root_ang_vel_w[0, 2]).item() > 0.1, (
            "Expected nonzero Z angular velocity from cross-product torque"
        )


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_global_force_at_position(device):
    """Global force at world position with non-identity rotation. Both rotation AND position correction matter."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(
            num_cubes=1, device=device, initial_rot=ROT_45_Z
        )

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Global force +X
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        forces[..., 0] = FORCE_MAGNITUDE
        torques = torch.zeros(1, len(body_ids), 3, device=device)

        # Position = each cube's link_pos + offset (same offset for both)
        offset = torch.zeros(1, len(body_ids), 3, device=device)
        offset[..., 1] = 1.0  # +1m Y offset in world frame

        pos_composer = cube_composer.data.body_com_pos_w[:, body_ids, :3].clone() + offset
        pos_raw = cube_raw.data.body_com_pos_w[:, body_ids, :3].clone() + offset

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            positions=pos_composer,
            body_ids=body_ids,
            is_global=True,
        )

        # Raw PhysX data
        raw_forces = torch.zeros(1, 3, device=device)
        raw_forces[:, 0] = FORCE_MAGNITUDE
        raw_torques = torch.zeros(1, 3, device=device)
        raw_positions = pos_raw.view(-1, 3)
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=raw_positions,
                indices=raw_indices,
                is_global=True,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Both linear and angular velocities should match
        torch.testing.assert_close(
            cube_composer.data.root_lin_vel_w,
            cube_raw.data.root_lin_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )
        torch.testing.assert_close(
            cube_composer.data.root_ang_vel_w,
            cube_raw.data.root_ang_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )

        # Sanity: angular velocity should be nonzero (cross-product torque)
        assert torch.abs(cube_composer.data.root_ang_vel_w[0, 2]).item() > 0.1, (
            "Expected nonzero Z angular velocity from positional torque"
        )


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_local_torque(device):
    """Local torque at identity orientation. Should produce matching angular velocity."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device)

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Composer path: local torque about +Z
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        torques = torch.zeros(1, len(body_ids), 3, device=device)
        torques[..., 2] = TORQUE_MAGNITUDE

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            body_ids=body_ids,
            is_global=False,
        )

        # Raw PhysX data
        raw_forces = torch.zeros(1, 3, device=device)
        raw_torques = torch.zeros(1, 3, device=device)
        raw_torques[:, 2] = TORQUE_MAGNITUDE
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=None,
                indices=raw_indices,
                is_global=False,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Angular velocities should match
        torch.testing.assert_close(
            cube_composer.data.root_ang_vel_w,
            cube_raw.data.root_ang_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )
        # Linear velocity should be ~zero for both (no force)
        torch.testing.assert_close(
            cube_composer.data.root_lin_vel_w,
            torch.zeros(1, 3, device=device),
            rtol=0.0,
            atol=1e-4,
        )
        torch.testing.assert_close(
            cube_raw.data.root_lin_vel_w,
            torch.zeros(1, 3, device=device),
            rtol=0.0,
            atol=1e-4,
        )


@pytest.mark.parametrize("device", ["cuda:0", "cpu"])
def test_composer_vs_physx_global_torque(device):
    """Global torque with non-identity rotation (45 deg Z). Composer rotates to body frame internally."""
    with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim:
        sim._app_control_on_stop_handle = None
        cube_composer, cube_raw = generate_dual_cube_scene(
            num_cubes=1, device=device, initial_rot=ROT_45_Z
        )

        sim.reset()

        body_ids, _ = cube_composer.find_bodies(".*")

        # Composer path: global torque about +Z
        forces = torch.zeros(1, len(body_ids), 3, device=device)
        torques = torch.zeros(1, len(body_ids), 3, device=device)
        torques[..., 2] = TORQUE_MAGNITUDE

        cube_composer.permanent_wrench_composer.set_forces_and_torques(
            forces=forces,
            torques=torques,
            body_ids=body_ids,
            is_global=True,
        )

        # Raw PhysX data
        raw_forces = torch.zeros(1, 3, device=device)
        raw_torques = torch.zeros(1, 3, device=device)
        raw_torques[:, 2] = TORQUE_MAGNITUDE
        raw_indices = cube_raw._ALL_INDICES

        for _ in range(N_STEPS):
            cube_composer.write_data_to_sim()
            cube_raw.write_data_to_sim()
            cube_raw.root_physx_view.apply_forces_and_torques_at_position(
                force_data=raw_forces,
                torque_data=raw_torques,
                position_data=None,
                indices=raw_indices,
                is_global=True,
            )
            sim.step()
            cube_composer.update(sim.cfg.dt)
            cube_raw.update(sim.cfg.dt)

        # Angular velocities should match
        torch.testing.assert_close(
            cube_composer.data.root_ang_vel_w,
            cube_raw.data.root_ang_vel_w,
            rtol=1e-4,
            atol=1e-4,
        )

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 18, 2026

I would try maybe with the branch before the merge of the composer to see how that goes? Misread that, that would be 2.3.1... That's weird.

@AntoineRichard
Copy link
Author

I did some more digging, and I can't find the bug. Could you share the mdp you're using to apply forces?

@AntoineRichard
Copy link
Author

AntoineRichard commented Feb 23, 2026

@LoreMoretti any chances you could share with me the mdp that causes problems :)

@LoreMoretti
Copy link

@AntoineRichard I am sorry for my very late reply. Let me check what I can share and I'll come back to you

@LoreMoretti
Copy link

Here the mdp:


def apply_payload(
    env: ManagerBasedEnv,
    env_ids: torch.Tensor,
    payload_range: list[tuple[float, float]],
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
    """Randomize the external forces and torques applied to the bodies.

    This function creates a set of random forces and torques sampled from the given ranges. The number of forces
    and torques is equal to the number of bodies times the number of environments. The forces and torques are
    applied to the bodies by calling ``asset.set_external_force_and_torque``. The forces and torques are only
    applied when ``asset.write_data_to_sim()`` is called in the environment.
    """
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject | Articulation = env.scene[asset_cfg.name]
    # resolve environment ids
    if env_ids is None:
        env_ids = torch.arange(env.scene.num_envs, device=asset.device)
    # resolve number of bodies
    num_bodies = (
        len(asset_cfg.body_ids)
        if isinstance(asset_cfg.body_ids, list)
        else asset.num_bodies
    )
    # sample random payloads
    size = (len(env_ids), num_bodies, 3)
    force_range = [
        tuple(p * 9.81 for p in pr) for pr in payload_range
    ]  # convert mass to force
    # forces is a tensor of shape (num_envs, num_bodies, 3), you should sample the z component only from the list of force_range
    forces = torch.zeros(size, device=asset.device)
    # Vectorized sampling: generate random values for all envs and bodies at once
    random_values = torch.rand((len(env_ids), num_bodies), device=asset.device)
    # Create tensors for min and max force values for each body
    force_min = torch.tensor([fr[0] for fr in force_range], device=asset.device)
    force_max = torch.tensor([fr[1] for fr in force_range], device=asset.device)
    # Scale random values to the force ranges
    forces[:, :, 2] = force_min + random_values * (force_max - force_min)
    torques = torch.zeros(size, device=asset.device)
    # set z towards the ground
    forces[:, :, 2] = -torch.abs(forces[:, :, 2])
    # set the forces and torques into the buffers
    # note: these are only applied when you call: `asset.write_data_to_sim()`
    asset.set_external_force_and_torque(
        forces, torques, env_ids=env_ids, body_ids=asset_cfg.body_ids, is_global=True
    )
    

Which is called at reset events.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants