1414
1515from __future__ import annotations
1616
17- import numpy as np
1817import torch
1918from typing import TYPE_CHECKING , Literal
2019
2524import omni .isaac .lab .utils .math as math_utils
2625from omni .isaac .lab .actuators import ImplicitActuator
2726from omni .isaac .lab .assets import Articulation , DeformableObject , RigidObject
28- from omni .isaac .lab .managers import SceneEntityCfg
27+ from omni .isaac .lab .managers import EventTermCfg , ManagerTermBase , SceneEntityCfg
2928from omni .isaac .lab .terrains import TerrainImporter
3029
3130if TYPE_CHECKING :
3231 from omni .isaac .lab .envs import ManagerBasedEnv
3332
3433
35- def randomize_rigid_body_material (
36- env : ManagerBasedEnv ,
37- env_ids : torch .Tensor | None ,
38- static_friction_range : tuple [float , float ],
39- dynamic_friction_range : tuple [float , float ],
40- restitution_range : tuple [float , float ],
41- num_buckets : int ,
42- asset_cfg : SceneEntityCfg ,
43- ):
34+ class randomize_rigid_body_material (ManagerTermBase ):
4435 """Randomize the physics materials on all geometries of the asset.
4536
4637 This function creates a set of physics materials with random static friction, dynamic friction, and restitution
@@ -53,76 +44,122 @@ def randomize_rigid_body_material(
5344 all bodies). The integer values are used as indices to select the material properties from the
5445 material buckets.
5546
47+ If the flag ``make_consistent`` is set to ``True``, the dynamic friction is set to be less than or equal to
48+ the static friction. This obeys the physics constraint on friction values. However, it may not always be
49+ essential for the application. Thus, the flag is set to ``False`` by default.
50+
5651 .. attention::
5752 This function uses CPU tensors to assign the material properties. It is recommended to use this function
5853 only during the initialization of the environment. Otherwise, it may lead to a significant performance
5954 overhead.
6055
6156 .. note::
6257 PhysX only allows 64000 unique physics materials in the scene. If the number of materials exceeds this
63- limit, the simulation will crash.
58+ limit, the simulation will crash. Due to this reason, we sample the materials only once during initialization.
59+ Afterwards, these materials are randomly assigned to the geometries of the asset.
6460 """
65- # extract the used quantities (to enable type-hinting)
66- asset : RigidObject | Articulation = env .scene [asset_cfg .name ]
6761
68- if not isinstance (asset , (RigidObject , Articulation )):
69- raise ValueError (
70- f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{ asset_cfg .name }
71- f" with type: '{ type (asset )}
72- )
62+ def __init__ (self , cfg : EventTermCfg , env : ManagerBasedEnv ):
63+ """Initialize the term.
7364
74- # resolve environment ids
75- if env_ids is None :
76- env_ids = torch .arange (env .scene .num_envs , device = "cpu" )
77- else :
78- env_ids = env_ids .cpu ()
79-
80- # retrieve material buffer
81- materials = asset .root_physx_view .get_material_properties ()
65+ Args:
66+ cfg: The configuration of the event term.
67+ env: The environment instance.
8268
83- # sample material properties from the given ranges
84- material_samples = np .zeros (materials [env_ids ].shape )
85- material_samples [..., 0 ] = np .random .uniform (* static_friction_range )
86- material_samples [..., 1 ] = np .random .uniform (* dynamic_friction_range )
87- material_samples [..., 2 ] = np .random .uniform (* restitution_range )
69+ Raises:
70+ ValueError: If the asset is not a RigidObject or an Articulation.
71+ """
72+ super ().__init__ (cfg , env )
8873
89- # create uniform range tensor for bucketing
90- lo = np . array ([ static_friction_range [ 0 ], dynamic_friction_range [ 0 ], restitution_range [ 0 ]])
91- hi = np . array ([ static_friction_range [ 1 ], dynamic_friction_range [ 1 ], restitution_range [ 1 ]])
74+ # extract the used quantities (to enable type-hinting)
75+ self . asset_cfg : SceneEntityCfg = cfg . params [ "asset_cfg" ]
76+ self . asset : RigidObject | Articulation = env . scene [ self . asset_cfg . name ]
9277
93- # to avoid 64k material limit in physx, we bucket materials by binning randomized material properties
94- # into buckets based on the number of buckets specified
95- for d in range ( 3 ):
96- buckets = np . array ([( hi [ d ] - lo [ d ]) * i / num_buckets + lo [ d ] for i in range ( num_buckets )])
97- material_samples [..., d ] = buckets [ np . searchsorted ( buckets , material_samples [..., d ]) - 1 ]
78+ if not isinstance ( self . asset , ( RigidObject , Articulation )):
79+ raise ValueError (
80+ f"Randomization term 'randomize_rigid_body_material' not supported for asset: ' { self . asset_cfg . name } '"
81+ f" with type: ' { type ( self . asset ) } '."
82+ )
9883
99- # update material buffer with new samples
100- if isinstance (asset , Articulation ) and asset_cfg .body_ids != slice (None ):
10184 # obtain number of shapes per body (needed for indexing the material properties correctly)
10285 # note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes
10386 # per body. We use the physics simulation view to obtain the number of shapes per body.
104- num_shapes_per_body = []
105- for link_path in asset .root_physx_view .link_paths [0 ]:
106- link_physx_view = asset ._physics_sim_view .create_rigid_body_view (link_path ) # type: ignore
107- num_shapes_per_body .append (link_physx_view .max_shapes )
87+ if isinstance (self .asset , Articulation ) and self .asset_cfg .body_ids != slice (None ):
88+ self .num_shapes_per_body = []
89+ for link_path in self .asset .root_physx_view .link_paths [0 ]:
90+ link_physx_view = self .asset ._physics_sim_view .create_rigid_body_view (link_path ) # type: ignore
91+ self .num_shapes_per_body .append (link_physx_view .max_shapes )
92+ # ensure the parsing is correct
93+ num_shapes = sum (self .num_shapes_per_body )
94+ expected_shapes = self .asset .root_physx_view .max_shapes
95+ if num_shapes != expected_shapes :
96+ raise ValueError (
97+ "Randomization term 'randomize_rigid_body_material' failed to parse the number of shapes per body."
98+ f" Expected total shapes: { expected_shapes } { num_shapes }
99+ )
100+ else :
101+ # in this case, we don't need to do special indexing
102+ self .num_shapes_per_body = None
103+
104+ # obtain parameters for sampling friction and restitution values
105+ static_friction_range = cfg .params .get ("static_friction_range" , (1.0 , 1.0 ))
106+ dynamic_friction_range = cfg .params .get ("dynamic_friction_range" , (1.0 , 1.0 ))
107+ restitution_range = cfg .params .get ("restitution_range" , (0.0 , 0.0 ))
108+ num_buckets = int (cfg .params .get ("num_buckets" , 1 ))
108109
109110 # sample material properties from the given ranges
110- for body_id in asset_cfg .body_ids :
111- # start index of shape
112- start_idx = sum (num_shapes_per_body [:body_id ])
113- # end index of shape
114- end_idx = start_idx + num_shapes_per_body [body_id ]
115- # assign the new materials
116- # material ids are of shape: num_env_ids x num_shapes
117- # material_buckets are of shape: num_buckets x 3
118- materials [env_ids , start_idx :end_idx ] = torch .from_numpy (material_samples [:, start_idx :end_idx ]).to (
119- dtype = torch .float
120- )
121- else :
122- materials [env_ids ] = torch .from_numpy (material_samples ).to (dtype = torch .float )
111+ # note: we only sample the materials once during initialization
112+ # afterwards these are randomly assigned to the geometries of the asset
113+ range_list = [static_friction_range , dynamic_friction_range , restitution_range ]
114+ ranges = torch .tensor (range_list , device = "cpu" )
115+ self .material_buckets = math_utils .sample_uniform (ranges [:, 0 ], ranges [:, 1 ], (num_buckets , 3 ), device = "cpu" )
116+
117+ # ensure dynamic friction is always less than static friction
118+ make_consistent = cfg .params .get ("make_consistent" , False )
119+ if make_consistent :
120+ self .material_buckets [:, 1 ] = torch .min (self .material_buckets [:, 0 ], self .material_buckets [:, 1 ])
121+
122+ def __call__ (
123+ self ,
124+ env : ManagerBasedEnv ,
125+ env_ids : torch .Tensor | None ,
126+ static_friction_range : tuple [float , float ],
127+ dynamic_friction_range : tuple [float , float ],
128+ restitution_range : tuple [float , float ],
129+ num_buckets : int ,
130+ asset_cfg : SceneEntityCfg ,
131+ make_consistent : bool = False ,
132+ ):
133+ # resolve environment ids
134+ if env_ids is None :
135+ env_ids = torch .arange (env .scene .num_envs , device = "cpu" )
136+ else :
137+ env_ids = env_ids .cpu ()
138+
139+ # randomly assign material IDs to the geometries
140+ total_num_shapes = self .asset .root_physx_view .max_shapes
141+ bucket_ids = torch .randint (0 , num_buckets , (len (env_ids ), total_num_shapes ), device = "cpu" )
142+ material_samples = self .material_buckets [bucket_ids ]
143+
144+ # retrieve material buffer from the physics simulation
145+ materials = self .asset .root_physx_view .get_material_properties ()
146+
147+ # update material buffer with new samples
148+ if self .num_shapes_per_body is not None :
149+ # sample material properties from the given ranges
150+ for body_id in self .asset_cfg .body_ids :
151+ # obtain indices of shapes for the body
152+ start_idx = sum (self .num_shapes_per_body [:body_id ])
153+ end_idx = start_idx + self .num_shapes_per_body [body_id ]
154+ # assign the new materials
155+ # material samples are of shape: num_env_ids x total_num_shapes x 3
156+ materials [env_ids , start_idx :end_idx ] = material_samples [:, start_idx :end_idx ]
157+ else :
158+ # assign all the materials
159+ materials [env_ids ] = material_samples [:]
123160
124- # apply to simulation
125- asset .root_physx_view .set_material_properties (materials , env_ids )
161+ # apply to simulation
162+ self . asset .root_physx_view .set_material_properties (materials , env_ids )
126163
127164
128165def randomize_rigid_body_mass (
0 commit comments