Skip to content

Commit 81490c2

Browse files
Milad-RakhshaMilad-Rakhsha
committed
Addressed the code review comments
Signed-off-by: Milad-Rakhsha <[email protected]>
1 parent c3220cb commit 81490c2

File tree

2 files changed

+103
-53
lines changed

2 files changed

+103
-53
lines changed

‎source/isaaclab/isaaclab/assets/articulation/articulation.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,23 @@ def root_newton_model(self) -> Model:
175175
"""Newton model for the asset."""
176176
return self._root_newton_view.model
177177

178+
@property
179+
def num_shapes_per_body(self) -> list[int]:
180+
"""Number of collision shapes per body in the articulation.
181+
182+
This property returns a list where each element represents the number of collision
183+
shapes for the corresponding body in the articulation. This is cached for efficient
184+
access during material property randomization and other operations.
185+
186+
Returns:
187+
List of integers representing the number of shapes per body.
188+
"""
189+
if not hasattr(self, "_num_shapes_per_body"):
190+
self._num_shapes_per_body = []
191+
for shapes in self._root_newton_view.body_shapes:
192+
self._num_shapes_per_body.append(len(shapes))
193+
return self._num_shapes_per_body
194+
178195
"""
179196
Operations.
180197
"""

‎source/isaaclab/isaaclab/envs/mdp/events.py‎

Lines changed: 86 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -64,34 +64,22 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
6464
f" with type: '{type(self.asset)}'."
6565
)
6666

67-
# obtain number of shapes per body (needed for indexing the material properties correctly)
68-
# note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes
69-
# per body. We use the physics simulation view to obtain the number of shapes per body.
70-
self.num_shapes_per_body = None
67+
# compute prefix scan for efficient indexing (shape counts come from Articulation)
68+
self.shape_start_indices = None
7169
if isinstance(self.asset, Articulation) and self.asset_cfg.body_ids != slice(None):
72-
self.num_shapes_per_body = []
73-
for shapes in self.asset.root_newton_view.body_shapes:
74-
self.num_shapes_per_body.append(len(shapes))
75-
76-
# obtain parameters for sampling friction and restitution values
77-
static_friction_range = cfg.params.get("static_friction_range", (1.0, 1.0))
78-
dynamic_friction_range = cfg.params.get("dynamic_friction_range", (1.0, 1.0))
79-
restitution_range = cfg.params.get("restitution_range", (0.0, 0.0))
80-
num_buckets = int(cfg.params.get("num_buckets", 1))
81-
82-
# sample material properties from the given ranges
83-
# note: we only sample the materials once during initialization
84-
# afterwards these are randomly assigned to the geometries of the asset
85-
range_list = [static_friction_range, dynamic_friction_range, restitution_range]
86-
ranges = torch.tensor(range_list, device=self.asset.device)
87-
self.material_buckets = math_utils.sample_uniform(
88-
ranges[:, 0], ranges[:, 1], (num_buckets, 3), device=self.asset.device
89-
)
70+
# get shapes per body from Articulation class
71+
num_shapes_per_body = self.asset.num_shapes_per_body
72+
73+
# compute prefix scan for faster lookup during randomization
74+
self.shape_start_indices = [0] # First body starts at index 0
75+
for i in range(len(num_shapes_per_body) - 1):
76+
self.shape_start_indices.append(self.shape_start_indices[-1] + num_shapes_per_body[i])
9077

91-
# ensure dynamic friction is always less than static friction
92-
make_consistent = cfg.params.get("make_consistent", False)
93-
if make_consistent:
94-
self.material_buckets[:, 1] = torch.min(self.material_buckets[:, 0], self.material_buckets[:, 1])
78+
# cache default material properties for consistent randomization baseline
79+
self.default_material_mu = None
80+
81+
# cache mask tensor for efficient reuse
82+
self.env_mask = torch.zeros((env.scene.num_envs,), dtype=torch.bool, device=env.device)
9583

9684
def __call__(
9785
self,
@@ -104,47 +92,92 @@ def __call__(
10492
asset_cfg: SceneEntityCfg,
10593
make_consistent: bool = False,
10694
):
95+
"""Randomize the material properties of rigid bodies.
96+
97+
This function randomizes the friction coefficient of rigid body shapes. Due to Newton physics
98+
(mjwarp) limitations, only the static friction coefficient is actually applied to the simulation.
99+
100+
Args:
101+
env: The environment instance.
102+
env_ids: The environment indices to apply the randomization to. If None, all environments are used.
103+
static_friction_range: Range for static friction coefficient. This is the only parameter that
104+
actually affects the simulation in Newton physics.
105+
dynamic_friction_range: Range for dynamic friction coefficient. **NOT USED** - Newton physics
106+
only supports a single friction coefficient per shape.
107+
restitution_range: Range for restitution coefficient. **NOT USED** - Newton physics does not
108+
currently support restitution randomization through this interface.
109+
num_buckets: Number of material buckets. **NOT USED** - materials are sampled dynamically.
110+
asset_cfg: The asset configuration specifying which asset and bodies to randomize.
111+
make_consistent: Whether to ensure dynamic friction <= static friction. **NOT USED** - only
112+
static friction is applied.
113+
114+
Note:
115+
Newton physics only supports setting a single friction coefficient (mu) per shape. The dynamic_friction_range and restitution_range parameters
116+
are kept for API consistency but do not affect the simulation.
117+
"""
107118

108119
# resolve environment ids
109120
if env_ids is None:
110121
env_ids = torch.arange(env.scene.num_envs, device="cpu")
111122
else:
112123
env_ids = env_ids.cpu()
113-
# retrieve material buffer from the physics simulation
114-
# materials = self.asset.root_physx_view.get_material_properties()
115-
material_mu = wp.to_torch(
116-
self.asset.root_newton_view.get_attribute("shape_material_mu", self.asset.root_newton_model)
117-
).clone()
118-
# randomly assign material IDs to the geometries
124+
125+
# env_ids is guaranteed to be a tensor at this point
126+
num_env_ids = env_ids.shape[0]
127+
128+
# cache default material properties on first call for consistent randomization baseline
129+
if self.default_material_mu is None:
130+
self.default_material_mu = wp.to_torch(
131+
self.asset.root_newton_view.get_attribute("shape_material_mu", self.asset.root_newton_model)
132+
).clone()
133+
134+
# start with default values and clone for safe modification
135+
material_mu = self.default_material_mu.clone()
136+
137+
# sample friction coefficients dynamically for each call
138+
# Newton physics uses a single friction coefficient (mu) per shape
139+
# Note: Only static_friction_range is used; dynamic_friction_range and restitution_range are ignored
119140
total_num_shapes = material_mu.shape[1]
120-
bucket_ids = torch.randint(0, num_buckets, (len(env_ids), total_num_shapes), device=self.asset.device)
121-
material_samples = self.material_buckets[bucket_ids]
122-
# print("material_mu before update:\n", material_mu.shape, material_mu)
123-
# update material buffer with new samples
124-
if self.num_shapes_per_body is not None:
125-
# sample material properties from the given ranges
141+
142+
# sample friction values directly for the environments and shapes that need updating
143+
if self.shape_start_indices is not None:
144+
# sample friction coefficients for specific body shapes using efficient indexing
126145
for body_id in self.asset_cfg.body_ids:
127-
# obtain indices of shapes for the body
128-
start_idx = sum(self.num_shapes_per_body[:body_id])
129-
end_idx = start_idx + self.num_shapes_per_body[body_id]
130-
# assign the new materials
131-
# material samples are of shape: num_env_ids x total_num_shapes x 3
132-
material_mu[env_ids, start_idx:end_idx] = material_samples[:, start_idx:end_idx, 0]
146+
# use precomputed indices for fast lookup
147+
start_idx = self.shape_start_indices[body_id]
148+
end_idx = start_idx + self.asset.num_shapes_per_body[body_id]
149+
num_shapes_in_body = end_idx - start_idx
150+
151+
# sample friction coefficients using only static_friction_range
152+
friction_samples = math_utils.sample_uniform(
153+
static_friction_range[0],
154+
static_friction_range[1],
155+
(num_env_ids, num_shapes_in_body),
156+
device=self.asset.device,
157+
)
158+
159+
# assign the new friction coefficients
160+
material_mu[env_ids, start_idx:end_idx] = friction_samples
133161
else:
134-
# assign all the materials
135-
material_mu[env_ids, :] = material_samples[:, :, 0]
162+
# sample friction coefficients for all shapes
163+
friction_samples = math_utils.sample_uniform(
164+
static_friction_range[0],
165+
static_friction_range[1],
166+
(num_env_ids, total_num_shapes),
167+
device=self.asset.device,
168+
)
136169

137-
# apply to simulation
138-
mask = torch.zeros((env.scene.num_envs,), dtype=torch.bool, device=env.device)
139-
mask[env_ids] = True
170+
# assign all the friction coefficients
171+
material_mu[env_ids, :] = friction_samples
172+
173+
# apply to simulation using cached mask
174+
self.env_mask.fill_(False) # reset all to False
175+
self.env_mask[env_ids] = True
140176
self.asset.root_newton_view.set_attribute(
141-
"shape_material_mu", self.asset.root_newton_model, wp.from_torch(material_mu), mask=mask
177+
"shape_material_mu", self.asset.root_newton_model, wp.from_torch(material_mu), mask=self.env_mask
142178
)
143179
NewtonManager._solver.notify_model_changed(SolverNotifyFlags.SHAPE_PROPERTIES)
144180

145-
# material_mu = wp.to_torch(self.asset.root_newton_view.get_attribute("shape_material_mu", self.asset.root_newton_model)).clone()
146-
# print("material_mu after update:\n", material_mu.shape, material_mu)
147-
148181

149182
def randomize_rigid_body_mass(
150183
env: ManagerBasedEnv,

0 commit comments

Comments
 (0)