1010import omni .log
1111from isaacsim .core .prims import XFormPrim
1212from isaacsim .core .utils .stage import get_current_stage
13- from pxr import Sdf , PhysxSchema
13+ from pxr import PhysxSchema , Sdf
1414
1515import isaaclab .sim as sim_utils
1616from isaaclab .assets import (
3131from isaaclab .sim .utils import get_current_stage_id
3232from isaaclab .terrains import TerrainImporter , TerrainImporterCfg
3333
34- from .cloner import physx_replicate , filter_collisions , usd_replicate , grid_transforms
34+ from .cloner import filter_collisions , grid_transforms , physx_replicate , usd_replicate
3535from .interactive_scene_cfg import InteractiveSceneCfg
3636
3737
@@ -145,16 +145,18 @@ def __init__(self, cfg: InteractiveSceneCfg):
145145 prim .CreateAttribute ("physxScene:envIdInBoundsBitCount" , Sdf .ValueTypeNames .Int ).Set (4 )
146146 self ._default_env_origins , _ = grid_transforms (self .cfg .num_envs , self .cfg .env_spacing , device = self .device )
147147 # copy empty prim of env_0 to env_1, env_2, ..., env_{num_envs-1} with correct location.
148- usd_replicate (self .stage , [self .env_fmt .format (0 )], [self .env_fmt ], self ._ALL_INDICES , positions = self ._default_env_origins )
148+ usd_replicate (
149+ self .stage , [self .env_fmt .format (0 )], [self .env_fmt ], self ._ALL_INDICES , positions = self ._default_env_origins
150+ )
149151
150152 self ._global_prim_paths = list ()
151153 if self ._is_scene_setup_from_cfg ():
152154 clone_plan = {"src" : [], "dest" : [], "mapping" : torch .empty ((0 ,), dtype = torch .bool ).to (self .device )}
153155 self ._global_template_prim_paths = list () # store paths that are in global collision filter from templates
154156 self ._add_entities_from_cfg ()
155157 for prim_path in self ._global_template_prim_paths :
156- i = clone_plan [' src' ].index (prim_path )
157- self ._global_prim_paths .extend (sim_utils .find_matching_prim_paths (clone_plan [' dest' ][i ].format (".*" )))
158+ i = clone_plan [" src" ].index (prim_path )
159+ self ._global_prim_paths .extend (sim_utils .find_matching_prim_paths (clone_plan [" dest" ][i ].format (".*" )))
158160
159161 self .clone_environments ()
160162
@@ -173,9 +175,9 @@ def clone_environments(self, copy_from_source: bool = False):
173175 clone_plan = {"src" : [], "dest" : [], "mapping" : torch .empty ((0 ,), dtype = torch .bool ).to (self .device )}
174176 prototypes = sim_utils .get_all_matching_child_prims (
175177 self .template_path ,
176- predicate = lambda prim : str (prim .GetPath ()).split ("/" )[- 1 ].startswith (self .prototype_name )
178+ predicate = lambda prim : str (prim .GetPath ()).split ("/" )[- 1 ].startswith (self .prototype_name ),
177179 )
178- prototype_root_set = set ([ "/" .join (str (prototype .GetPath ()).split ("/" )[:- 1 ]) for prototype in prototypes ])
180+ prototype_root_set = { "/" .join (str (prototype .GetPath ()).split ("/" )[:- 1 ]) for prototype in prototypes }
179181 for prototype_root in prototype_root_set :
180182 protos = sim_utils .find_matching_prim_paths (f"{ prototype_root } /.*" )
181183 protos = [proto for proto in protos if proto .split ("/" )[- 1 ].startswith (self .prototype_name )]
@@ -187,22 +189,22 @@ def clone_environments(self, copy_from_source: bool = False):
187189 else :
188190 m [self ._ALL_INDICES % len (protos ), self ._ALL_INDICES ] = True
189191
190- clone_plan [' src' ].extend (protos )
191- clone_plan [' dest' ].extend ([prototype_root .replace (self .template_path , self .env_fmt )] * len (protos ))
192- clone_plan [' mapping' ] = torch .cat ((clone_plan [' mapping' ].reshape (- 1 , m .size (1 )), m ), dim = 0 )
192+ clone_plan [" src" ].extend (protos )
193+ clone_plan [" dest" ].extend ([prototype_root .replace (self .template_path , self .env_fmt )] * len (protos ))
194+ clone_plan [" mapping" ] = torch .cat ((clone_plan [" mapping" ].reshape (- 1 , m .size (1 )), m ), dim = 0 )
193195
194- proto_idx = clone_plan [' mapping' ].to (torch .int32 ).argmax (dim = 1 )
195- proto_mask = torch .zeros_like (clone_plan [' mapping' ])
196- proto_mask .scatter_ (1 , proto_idx .view (- 1 , 1 ).to (torch .long ), clone_plan [' mapping' ].any (dim = 1 , keepdim = True ))
197- usd_replicate (self .stage , clone_plan [' src' ], clone_plan [' dest' ], self ._ALL_INDICES , proto_mask )
196+ proto_idx = clone_plan [" mapping" ].to (torch .int32 ).argmax (dim = 1 )
197+ proto_mask = torch .zeros_like (clone_plan [" mapping" ])
198+ proto_mask .scatter_ (1 , proto_idx .view (- 1 , 1 ).to (torch .long ), clone_plan [" mapping" ].any (dim = 1 , keepdim = True ))
199+ usd_replicate (self .stage , clone_plan [" src" ], clone_plan [" dest" ], self ._ALL_INDICES , proto_mask )
198200 self .stage .GetPrimAtPath (self .template_path ).SetActive (False )
199201
200202 # If all prototypes map to env_0, clone whole env_0 to all envs; else clone per-object
201203 if torch .all (proto_idx == 0 ):
202- replicate_args = [self .env_fmt .format (0 )], [self .env_fmt ], self ._ALL_INDICES , clone_plan [' mapping' ]
204+ replicate_args = [self .env_fmt .format (0 )], [self .env_fmt ], self ._ALL_INDICES , clone_plan [" mapping" ]
203205 else :
204- src = [tpl .format (int (idx )) for tpl , idx in zip (clone_plan [' dest' ], proto_idx .tolist ())]
205- replicate_args = src , clone_plan [' dest' ], self ._ALL_INDICES , clone_plan [' mapping' ]
206+ src = [tpl .format (int (idx )) for tpl , idx in zip (clone_plan [" dest" ], proto_idx .tolist ())]
207+ replicate_args = src , clone_plan [" dest" ], self ._ALL_INDICES , clone_plan [" mapping" ]
206208 else :
207209 mapping = torch .ones ((1 , self .num_envs ), device = self .device , dtype = torch .bool )
208210 replicate_args = [self .env_fmt .format (0 )], [self .env_fmt ], self ._ALL_INDICES , mapping
@@ -677,7 +679,7 @@ def _add_entities_from_cfg(self):
677679 # prepared /World/template path, once all template is ready, cloner can determine what rules to follow
678680 # to combine, and distribute the templates to cloned environments.
679681 destinations_regex_ns = asset_cfg .prim_path .format (ENV_REGEX_NS = self .env_regex_ns )
680- if ( self .env_regex_ns [:- 2 ] in destinations_regex_ns ) :
682+ if self .env_regex_ns [:- 2 ] in destinations_regex_ns :
681683 require_clone = True
682684 prototype_root = asset_cfg .prim_path .format (ENV_REGEX_NS = self .template_path )
683685 asset_cfg .prim_path = f"{ prototype_root } /{ self .prototype_name } _.*"
0 commit comments