diff --git a/omnigibson/envs/__init__.py b/omnigibson/envs/__init__.py index 8cc3c969a..a93206368 100644 --- a/omnigibson/envs/__init__.py +++ b/omnigibson/envs/__init__.py @@ -1 +1,2 @@ from omnigibson.envs.env_base import Environment +from omnigibson.envs.env_wrapper import EnvironmentWrapper, create_wrapper, REGISTERED_ENV_WRAPPERS diff --git a/omnigibson/envs/env_base.py b/omnigibson/envs/env_base.py index f02f83333..7ef4a4578 100644 --- a/omnigibson/envs/env_base.py +++ b/omnigibson/envs/env_base.py @@ -1,5 +1,6 @@ import gym import numpy as np +from copy import deepcopy import omnigibson as og from omnigibson.objects import REGISTERED_OBJECTS @@ -50,7 +51,9 @@ def __init__( self._automatic_reset = automatic_reset self._flatten_action_space = flatten_action_space self._flatten_obs_space = flatten_obs_space + self.physics_timestep = physics_timestep self.action_timestep = action_timestep + self.device = device # Initialize other placeholders that will be filled in later self._initial_pos_z_offset = None # how high to offset object placement to account for one action step of dropping @@ -71,12 +74,6 @@ def __init__( for config in configs: merge_nested_dicts(base_dict=self.config, extra_dict=parse_config(config), inplace=True) - # Set the simulator settings - og.sim.set_simulation_dt(physics_dt=physics_timestep, rendering_dt=action_timestep) - og.sim.viewer_width = self.render_config["viewer_width"] - og.sim.viewer_height = self.render_config["viewer_height"] - og.sim.device = device - # Load this environment self.load() @@ -153,10 +150,22 @@ def _load_variables(self): drop_distance = 0.5 * 9.8 * (self.action_timestep ** 2) assert drop_distance < self._initial_pos_z_offset, "initial_pos_z_offset is too small for collision checking" - def _load_task(self): + def _load_task(self, task_config=None): """ Load task + + Args: + task_confg (None or dict): If specified, custom task configuration to use. Otherwise, will use + self.task_config. Note that if a custom task configuration is specified, the internal task config + will be updated as well """ + # Update internal config if specified + if task_config is not None: + # Copy task config, in case self.task_config and task_config are the same! + task_config = deepcopy(task_config) + self.task_config.clear() + self.task_config.update(task_config) + # Sanity check task to make sure it's valid task_type = self.task_config["type"] assert_valid_key(key=task_type, valid_keys=REGISTERED_TASKS, name="task type") @@ -188,6 +197,13 @@ def _load_scene(self): cls_type_descriptor="scene", ) og.sim.import_scene(scene) + + # Set the simulator settings + og.sim.set_simulation_dt(physics_dt=self.physics_timestep, rendering_dt=self.action_timestep) + og.sim.viewer_width = self.render_config["viewer_width"] + og.sim.viewer_height = self.render_config["viewer_height"] + og.sim.device = self.device + assert og.sim.is_stopped(), "Simulator must be stopped after loading scene!" def _load_robots(self): @@ -324,6 +340,31 @@ def load(self): # Denote that the scene is loaded self._loaded = True + def update_task(self, task_config): + """ + Updates the internal task using @task_config. NOTE: This will internally reset the environment as well! + + Args: + task_config (dict): Task configuration for updating the new task + """ + # Make sure sim is playing + assert og.sim.is_playing(), "Update task should occur while sim is playing!" + + # Denote scene as not loaded yet + self._loaded = False + og.sim.stop() + self._load_task(task_config=task_config) + og.sim.play() + self.reset() + + # Load obs / action spaces + self.load_observation_space() + self._load_action_space() + + # Scene is now loaded again + self._loaded = True + + def close(self): """ Clean up the environment and shut down the simulation. @@ -443,7 +484,7 @@ def reset(self): # Grab and return observations obs = self.get_obs() - if self.observation_space is not None and not self.observation_space.contains(obs): + if self._loaded and not self.observation_space.contains(obs): # Flatten obs, and print out all keys and values log.error("OBSERVATION SPACE:") for key, value in recursively_generate_flat_dict(dic=self.observation_space).items(): @@ -543,6 +584,14 @@ def task_config(self): """ return self.config["task"] + @property + def wrapper_config(self): + """ + Returns: + dict: Wrapper-specific configuration kwargs + """ + return self.config["wrapper"] + @property def default_config(self): """ @@ -584,5 +633,10 @@ def default_config(self): # Task kwargs "task": { "type": "DummyTask", - } + }, + + # Wrapper kwargs + "wrapper": { + "type": None, + }, } diff --git a/omnigibson/envs/env_wrapper.py b/omnigibson/envs/env_wrapper.py index e49a96fa7..0561d6da8 100644 --- a/omnigibson/envs/env_wrapper.py +++ b/omnigibson/envs/env_wrapper.py @@ -1,7 +1,32 @@ from omnigibson.utils.python_utils import Wrapper +from omnigibson.utils.python_utils import Registerable, classproperty, create_class_from_registry_and_config +from omnigibson.utils.ui_utils import create_module_logger +from copy import deepcopy +# Global dicts that will contain mappings +REGISTERED_ENV_WRAPPERS = dict() -class EnvironmentWrapper(Wrapper): +# Create module logger +log = create_module_logger(module_name=__name__) + + +def create_wrapper(env): + """ + Wraps environment @env with wrapper defined by env.wrapper_config + """ + wrapper_cfg = deepcopy(env.wrapper_config) + wrapper_type = wrapper_cfg.pop("type") + wrapper_cfg["env"] = env + + return create_class_from_registry_and_config( + cls_name=wrapper_type, + cls_registry=REGISTERED_ENV_WRAPPERS, + cfg=wrapper_cfg, + cls_type_descriptor="wrapper", + ) + + +class EnvironmentWrapper(Wrapper, Registerable): """ Base class for all environment wrappers in OmniGibson. In general, reset(), step(), and observation_spec() should be overwritten @@ -50,3 +75,15 @@ def observation_spec(self): """ return self.env.observation_spec() + @classproperty + def _do_not_register_classes(cls): + # Don't register this class since it's an abstract template + classes = super()._do_not_register_classes + classes.add("EnvironmentWrapper") + return classes + + @classproperty + def _cls_registry(cls): + # Global robot registry + global REGISTERED_ENV_WRAPPERS + return REGISTERED_ENV_WRAPPERS diff --git a/omnigibson/tasks/behavior_task.py b/omnigibson/tasks/behavior_task.py index 395e30361..3b886895a 100644 --- a/omnigibson/tasks/behavior_task.py +++ b/omnigibson/tasks/behavior_task.py @@ -17,6 +17,7 @@ from omnigibson.robots.robot_base import BaseRobot from omnigibson.systems.system_base import get_system, add_callback_on_system_init, add_callback_on_system_clear, \ REGISTERED_SYSTEMS +from omnigibson.scenes.scene_base import Scene from omnigibson.scenes.interactive_traversable_scene import InteractiveTraversableScene from omnigibson.utils.bddl_utils import OmniGibsonBDDLBackend, BDDLEntity, BEHAVIOR_ACTIVITIES, BDDLSampler from omnigibson.tasks.task_base import BaseTask @@ -142,7 +143,7 @@ def verify_scene_and_task_config(cls, scene_cfg, task_cfg): task_cfg.get("predefined_problem", None) is not None else task_cfg["activity_name"] if scene_file is None and scene_instance is None and not task_cfg["online_object_sampling"]: scene_instance = cls.get_cached_activity_scene_filename( - scene_model=scene_cfg["scene_model"], + scene_model=scene_cfg.get("scene_model", "Scene"), activity_name=activity_name, activity_definition_id=task_cfg.get("activity_definition_id", 0), activity_instance_id=task_cfg.get("activity_instance_id", 0), @@ -522,8 +523,8 @@ def name(self): @classproperty def valid_scene_types(cls): - # Must be an interactive traversable scene - return {InteractiveTraversableScene} + # Any scene can be used + return {Scene} @classproperty def default_termination_config(cls): diff --git a/omnigibson/tasks/task_base.py b/omnigibson/tasks/task_base.py index 6ec84aeb3..46d3eed04 100644 --- a/omnigibson/tasks/task_base.py +++ b/omnigibson/tasks/task_base.py @@ -50,6 +50,7 @@ def __init__(self, termination_config=None, reward_config=None): self._loaded = False self._reward = None self._done = None + self._success = None self._info = None self._low_dim_obs_dim = None @@ -160,7 +161,8 @@ def _reset_variables(self, env): """ # By default, reset reward, done, and info self._reward = None - self._done = None + self._done = False + self._success = False self._info = None def reset(self, env): @@ -311,6 +313,7 @@ def step(self, env, action): # Update the internal state of this task self._reward = reward self._done = done + self._success = done_info["success"] self._info = { "reward": reward_info, "done": done_info, @@ -344,6 +347,15 @@ def done(self): assert self._done is not None, "At least one step() must occur before done can be calculated!" return self._done + @property + def success(self): + """ + Returns: + bool: Whether this task has succeeded or not + """ + assert self._success is not None, "At least one step() must occur before success can be calculated!" + return self._success + @property def info(self): """ diff --git a/omnigibson/utils/bddl_utils.py b/omnigibson/utils/bddl_utils.py index c80f22b25..f7493e271 100644 --- a/omnigibson/utils/bddl_utils.py +++ b/omnigibson/utils/bddl_utils.py @@ -24,6 +24,7 @@ from omnigibson import object_states from omnigibson.object_states.factory import _KINEMATIC_STATE_SET from omnigibson.systems.system_base import is_system_active, get_system +from omnigibson.scenes.interactive_traversable_scene import InteractiveTraversableScene # Create module logger log = create_module_logger(module_name=__name__) @@ -276,7 +277,7 @@ def __init__( ): # Store internal variables from inputs self._env = env - self._scene_model = self._env.scene.scene_model + self._scene_model = self._env.scene.scene_model if isinstance(self._env.scene, InteractiveTraversableScene) else None self._agent = self._env.robots[0] if debug: gm.DEBUG = True @@ -411,9 +412,9 @@ def _parse_inroom_object_room_assignment(self): # Invalid room assignment return f"You have assigned room type for [{obj_synset}], but [{obj_synset}] is sampleable. " \ f"Only non-sampleable (scene) objects can have room assignment." - if room_type not in og.sim.scene.seg_map.room_sem_name_to_ins_name: + if self._scene_model is not None and room_type not in og.sim.scene.seg_map.room_sem_name_to_ins_name: # Missing room type - return f"Room type [{room_type}] missing in scene [{og.sim.scene.scene_model}]." + return f"Room type [{room_type}] missing in scene [{self._scene_model}]." if room_type not in self._room_type_to_object_instance: self._room_type_to_object_instance[room_type] = [] self._room_type_to_object_instance[room_type].append(obj_inst) @@ -521,6 +522,12 @@ def _build_sampling_order(self): # Sanity check kinematic objects -- any non-system must be kinematically sampled remaining_kinematic_entities = nonparticle_entities - unsampleable_obj_instances - \ self._inroom_object_instances - set.union(*(self._object_sampling_orders["kinematic"] + [set()])) + + # Possibly remove the agent entity if we're in an empty scene -- i.e.: no kinematic sampling needed for the + # agent + if self._scene_model is None: + remaining_kinematic_entities -= {"agent.n.01_1"} + if len(remaining_kinematic_entities) != 0: return f"Some objects do not have any kinematic condition defined for them in the initial conditions: " \ f"{', '.join(remaining_kinematic_entities)}" @@ -566,7 +573,8 @@ def _build_inroom_object_scope(self): valid_models = {cat: set(get_all_object_category_models_with_abilities(cat, abilities)) for cat in categories} - for room_inst in og.sim.scene.seg_map.room_sem_name_to_ins_name[room_type]: + room_insts = [None] if self._scene_model is None else og.sim.scene.seg_map.room_sem_name_to_ins_name[room_type] + for room_inst in room_insts: # A list of scene objects that satisfy the requested categories room_objs = og.sim.scene.object_registry("in_rooms", room_inst, default_val=[]) scene_objs = [obj for obj in room_objs if obj.category in categories and obj.model in valid_models[obj.category]] @@ -660,11 +668,15 @@ def _filter_object_scope(self, input_object_scope, conditions, condition_type): filtered_object_scope[room_type][scene_obj][room_inst].append(obj) # Compute most problematic objects - problematic_objs_by_proportion = defaultdict(list) - for child_scope_name, parent_obj_names in problematic_objs.items(): - problematic_objs_by_proportion[np.mean(list(parent_obj_names.values()))].append(child_scope_name) + if len(problematic_objs) == 0: + max_problematic_objs = [] + else: + problematic_objs_by_proportion = defaultdict(list) + for child_scope_name, parent_obj_names in problematic_objs.items(): + problematic_objs_by_proportion[np.mean(list(parent_obj_names.values()))].append(child_scope_name) + max_problematic_objs = problematic_objs_by_proportion[min(problematic_objs_by_proportion.keys())] - return filtered_object_scope, problematic_objs_by_proportion[min(problematic_objs_by_proportion.keys())] + return filtered_object_scope, max_problematic_objs def _consolidate_room_instance(self, filtered_object_scope, condition_type): """