Makes BaseEnv and RLEnv consistent with desired definitions (#197)

# Description

As discussed in the past, the `BaseEnv` should only contain an
observation-action interface, while `RlEnv` should contain also the
task-specific interface. Based on this, the MR introduces the following
changes:

* Moves randomization manager to the `BaseEnv` class with the default
settings to reset the scene to the defaults specified in the
configurations of assets.
* Moves command generator to `RlEnv` class
* Adds a `NullCommandGenerator` for no command environments: This is
easier to work with than having checks for None.

## Type of change

- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

---------

Signed-off-by: Mayank Mittal <[email protected]>
Co-authored-by: jsmith-bdai <[email protected]>

source/extensions/omni.isaac.orbit/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.21"
+version = "0.9.22"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

@@ -1,6 +1,24 @@
 Changelog
 ---------
 
+0.9.22 (2023-10-26)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added a :class:`omni.isaac.orbit.command_generators.NullCommandGenerator` class for no command environments.
+  This is easier to work with than having checks for :obj:`None` in the command generator.
+
+Fixed
+^^^^^
+
+* Moved the randomization manager to the :class:`omni.isaac.orbit.envs.BaseEnv` class with the default
+  settings to reset the scene to the defaults specified in the configurations of assets.
+* Moved command generator to the :class:`omni.isaac.orbit.envs.RlEnv` class to have all task-specification
+  related classes in the same place.
+
+
 0.9.21 (2023-10-26)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/__init__.py

@@ -21,15 +21,20 @@
 from .command_generator_cfg import (
     CommandGeneratorBaseCfg,
     NormalVelocityCommandGeneratorCfg,
+    NullCommandGeneratorCfg,
     TerrainBasedPositionCommandGeneratorCfg,
     UniformVelocityCommandGeneratorCfg,
 )
+from .null_command_generator import NullCommandGenerator
 from .position_command_generator import TerrainBasedPositionCommandGenerator
 from .velocity_command_generator import NormalVelocityCommandGenerator, UniformVelocityCommandGenerator
 
 __all__ = [
     "CommandGeneratorBase",
     "CommandGeneratorBaseCfg",
+    # null command generator
+    "NullCommandGenerator",
+    "NullCommandGeneratorCfg",
     # velocity command generators
     "UniformVelocityCommandGenerator",
     "UniformVelocityCommandGeneratorCfg",

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_cfg.py

@@ -5,11 +5,13 @@
 
 from __future__ import annotations
 
+import math
 from dataclasses import MISSING
 
 from omni.isaac.orbit.utils import configclass
 
 from .command_generator_base import CommandGeneratorBase
+from .null_command_generator import NullCommandGenerator
 from .position_command_generator import TerrainBasedPositionCommandGenerator
 from .velocity_command_generator import NormalVelocityCommandGenerator, UniformVelocityCommandGenerator
 
@@ -34,6 +36,23 @@ class CommandGeneratorBaseCfg:
     """Whether to visualize debug information. Defaults to False."""
 
 
+"""
+Null-command generator.
+"""
+
+
+@configclass
+class NullCommandGeneratorCfg(CommandGeneratorBaseCfg):
+    """Configuration for the null command generator."""
+
+    class_type: type = NullCommandGenerator
+
+    def __post_init__(self):
+        """Post initialization."""
+        # set the resampling time range to infinity to avoid resampling
+        self.resampling_time_range = (math.inf, math.inf)
+
+
 """
 Locomotion-specific command generators.
 """

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/null_command_generator.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Sub-module containing command generator that does nothing."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from omni.isaac.orbit.command_generators.command_generator_base import CommandGeneratorBase
+
+if TYPE_CHECKING:
+    from .command_generator_cfg import NullCommandGeneratorCfg
+
+
+class NullCommandGenerator(CommandGeneratorBase):
+    """Command generator that does nothing.
+
+    This command generator does not generate any commands. It is used for environments that do not
+    require any commands.
+    """
+
+    cfg: NullCommandGeneratorCfg
+    """Configuration for the command generator."""
+
+    """
+    Properties
+    """
+
+    @property
+    def command(self):
+        """Null command.
+
+        Raises:
+            RuntimeError: No command is generated. Always raises this error.
+        """
+        raise RuntimeError("NullCommandGenerator does not generate any commands.")
+
+    """
+    Implementation specific functions.
+    """
+
+    def _resample_command(self, env_ids: Sequence[int]):
+        pass
+
+    def _update_command(self):
+        pass
+
+    def _update_metrics(self):
+        pass

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/__init__.py

@@ -5,16 +5,18 @@
 
 from .base_env import BaseEnv
 from .base_env_cfg import BaseEnvCfg, ViewerCfg
-from .rl_env import RLEnv, VecEnvIndices, VecEnvObs, VecEnvStepReturn
+from .rl_env import RLEnv, VecEnvObs, VecEnvStepReturn
 from .rl_env_cfg import RLEnvCfg
 
 __all__ = [
+    # base
     "BaseEnv",
     "BaseEnvCfg",
     "ViewerCfg",
+    # rl
     "RLEnv",
     "RLEnvCfg",
-    "VecEnvIndices",
+    # env type variables
     "VecEnvObs",
     "VecEnvStepReturn",
 ]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env.py

@@ -7,8 +7,9 @@
 
 import builtins
 
-from omni.isaac.orbit.command_generators import CommandGeneratorBase
-from omni.isaac.orbit.managers import ActionManager, ObservationManager
+import omni.isaac.core.utils.torch as torch_utils
+
+from omni.isaac.orbit.managers import ActionManager, ObservationManager, RandomizationManager
 from omni.isaac.orbit.scene import InteractiveScene
 from omni.isaac.orbit.sim import SimulationContext
 from omni.isaac.orbit.utils.timer import Timer
@@ -36,10 +37,10 @@ class BaseEnv:
       raw actions at different levels of abstraction. For example, in case of a robotic arm, the raw actions
       can be joint torques, joint positions, or end-effector poses. Similarly for a mobile base, it can be
       the joint torques, or the desired velocity of the floating base.
-    * **Command Generator**: The command generator that generates the goal commands for the robot. These
-      commands are used by the observation manager to generate the observations. For example, in case of a
-      robotic arm, the goal commands can be the object to be grasped, or the desired end-effector pose. For
-      a mobile base, it can be the goal position and orientation of the base.
+    * **Randomization Manager**: The randomization manager that randomizes different elements in the scene.
+      This includes resetting the scene to a default state or randomize the scene at different intervals
+      of time. The randomization manager can be configured to randomize different elements of the scene
+      such as the masses of objects, friction coefficients, or apply random pushes to the robot.
 
     The environment provides a unified interface for interacting with the simulation. However, it does not
     include task-specific quantities such as the reward function, or the termination conditions. These
@@ -135,7 +136,7 @@ def device(self):
         return self.sim.device
 
     """
-    Operations.
+    Operations - Setup.
     """
 
     def load_managers(self):
@@ -145,16 +146,34 @@ def load_managers(self):
             This must happen after the simulator is reset, i.e. after the first call to :meth:`self.sim.reset`.
         """
         # prepare the managers
-        # note: this order is important since observation manager needs to know the command and action managers
-        # -- command manager
-        self.command_manager: CommandGeneratorBase = self.cfg.commands.class_type(self.cfg.commands, self)
-        print("[INFO] Command Manager: ", self.command_manager)
         # -- action manager
         self.action_manager = ActionManager(self.cfg.actions, self)
         print("[INFO] Action Manager: ", self.action_manager)
         # -- observation manager
         self.observation_manager = ObservationManager(self.cfg.observations, self)
         print("[INFO] Observation Manager:", self.observation_manager)
+        # -- randomization manager
+        self.randomization_manager = RandomizationManager(self.cfg.randomization, self)
+        print("[INFO] Randomization Manager: ", self.randomization_manager)
+
+    """
+    Operations - MDP.
+    """
+
+    @staticmethod
+    def seed(seed: int = -1) -> int:
+        """Set the seed for the environment.
+
+        Args:
+            seed: The seed for random generator. Defaults to -1.
+
+        Returns:
+            The seed used for random generator.
+        """
+        import omni.replicator.core as rep
+
+        rep.set_global_seed(seed)
+        return torch_utils.set_seed(seed)
 
     def close(self):
         """Cleanup for the environment."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env_cfg.py

@@ -12,7 +12,8 @@
 from dataclasses import MISSING
 from typing import Tuple
 
-from omni.isaac.orbit.command_generators import CommandGeneratorBaseCfg
+import omni.isaac.orbit.envs.mdp as mdp
+from omni.isaac.orbit.managers import RandomizationTermCfg as RandTerm
 from omni.isaac.orbit.scene import InteractiveSceneCfg
 from omni.isaac.orbit.sim import SimulationCfg
 from omni.isaac.orbit.utils import configclass
@@ -38,6 +39,17 @@ class ViewerCfg:
     """
 
 
+@configclass
+class DefaultRandomizationManagerCfg:
+    """Configuration of the default randomization manager.
+
+    This manager is used to reset the scene to a default state. The default state is specified
+    by the scene configuration.
+    """
+
+    reset_scene_to_default = RandTerm(func=mdp.reset_scene_to_default, mode="reset")
+
+
 @configclass
 class BaseEnvCfg:
     """Base configuration of the environment."""
@@ -59,5 +71,6 @@ class BaseEnvCfg:
     """Observation space settings."""
     actions: object = MISSING
     """Action space settings."""
-    commands: CommandGeneratorBaseCfg = MISSING
-    """Command generator settings."""
+    randomization: object = DefaultRandomizationManagerCfg()
+    """Randomization settings. Defaults to the default randomization manager, which resets
+    the scene to its default state."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/randomizations.py

@@ -22,11 +22,11 @@
 from omni.isaac.orbit.utils.math import quat_from_euler_xyz, sample_uniform
 
 if TYPE_CHECKING:
-    from omni.isaac.orbit.envs.rl_env import RLEnv
+    from omni.isaac.orbit.envs.base_env import BaseEnv
 
 
 def randomize_rigid_body_material(
-    env: RLEnv,
+    env: BaseEnv,
     env_ids: torch.Tensor | None,
     static_friction_range: tuple[float, float],
     dynamic_friction_range: tuple[float, float],
@@ -79,7 +79,9 @@ def randomize_rigid_body_material(
     asset.body_physx_view.set_material_properties(materials, indices)
 
 
-def add_body_mass(env: RLEnv, env_ids: torch.Tensor | None, mass_range: tuple[float, float], asset_cfg: SceneEntityCfg):
+def add_body_mass(
+    env: BaseEnv, env_ids: torch.Tensor | None, mass_range: tuple[float, float], asset_cfg: SceneEntityCfg
+):
     """Randomize the mass of the bodies by adding a random value sampled from the given range.
 
     .. tip::
@@ -107,7 +109,7 @@ def add_body_mass(env: RLEnv, env_ids: torch.Tensor | None, mass_range: tuple[fl
 
 
 def apply_external_force_torque(
-    env: RLEnv,
+    env: BaseEnv,
     env_ids: torch.Tensor,
     force_range: tuple[float, float],
     torque_range: tuple[float, float],
@@ -137,7 +139,7 @@ def apply_external_force_torque(
 
 
 def push_by_setting_velocity(
-    env: RLEnv,
+    env: BaseEnv,
     env_ids: torch.Tensor,
     velocity_range: dict[str, tuple[float, float]],
     asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
@@ -168,7 +170,7 @@ def push_by_setting_velocity(
 
 
 def reset_root_state(
-    env: RLEnv,
+    env: BaseEnv,
     env_ids: torch.Tensor,
     pose_range: dict[str, tuple[float, float]],
     velocity_range: dict[str, tuple[float, float]],
@@ -219,7 +221,7 @@ def reset_root_state(
 
 
 def reset_joints_by_scale(
-    env: RLEnv,
+    env: BaseEnv,
     env_ids: torch.Tensor,
     position_range: tuple[float, float],
     velocity_range: tuple[float, float],
@@ -241,3 +243,21 @@ def reset_joints_by_scale(
 
     # set into the physics simulation
     asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
+
+
+def reset_scene_to_default(env: BaseEnv, env_ids: torch.Tensor):
+    """Reset the scene to the default state specified in the scene configuration."""
+    # root states
+    for rigid_object in env.scene.rigid_objects.values() + env.scene.articulations.values():
+        # obtain default and deal with the offset for env origins
+        default_root_state = rigid_object.data.default_root_state_w[env_ids].clone()
+        default_root_state[:, 0:3] += env.scene.env_origins[env_ids]
+        # set into the physics simulation
+        rigid_object.write_root_state_to_sim(default_root_state, env_ids=env_ids)
+    # joint states
+    for articulation_asset in env.scene.articulations.values():
+        # obtain default joint positions
+        default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone()
+        default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone()
+        # set into the physics simulation
+        articulation_asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids)