@@ -56,7 +56,9 @@ def __init__(
5656 satellites : Union [Satellite , list [Satellite ]],
5757 scenario : Optional [Scenario ] = None ,
5858 rewarder : Optional [Union [GlobalReward , list [GlobalReward ]]] = None ,
59- world_type : Optional [type [WorldModel ]] = None ,
59+ world_type : Optional [
60+ Union [type [WorldModel ], tuple [type [WorldModel ], ...]]
61+ ] = None ,
6062 world_args : Optional [dict [str , Any ]] = None ,
6163 communicator : Optional [CommunicationMethod ] = None ,
6264 sat_arg_randomizer : Optional [SatArgRandomizer ] = None ,
@@ -96,7 +98,7 @@ def __init__(
9698 sat_arg_randomizer: For correlated randomization of satellites arguments. Should
9799 be a function that takes a list of satellites and returns a dictionary that
98100 maps satellites to dictionaries of satellite model arguments to be overridden.
99- world_type: Type of Basilisk world model to be constructed.
101+ world_type: Type or tuple of types of Basilisk world model to be constructed.
100102 world_args: Arguments for :class:`~bsk_rl.sim.world.WorldModel` construction.
101103 Should be in the form of a dictionary with keys corresponding to the
102104 arguments of the constructor and values that are either the desired value
@@ -175,6 +177,14 @@ def __init__(
175177
176178 if world_type is None :
177179 world_type = self ._minimum_world_model ()
180+ else :
181+ logger .warning (
182+ "Using user-specified world type. Generally, the env-determined world "
183+ "type is sufficient."
184+ )
185+ if not isinstance (world_type , (list , tuple )):
186+ world_type = (world_type ,)
187+ world_type = functional .compose_types (* world_type , WorldModel , name = "World" )
178188 self .world_type = world_type
179189 if world_args is None :
180190 world_args = self .world_type .default_world_args ()
0 commit comments