Train High-Level Policies in Hierarchical Approaches

habitat-baselines/habitat_baselines/README.md

@@ -53,6 +53,31 @@ To use them download pre-trained pytorch models from [link](https://dl.fbaipubli
 The `habitat_baselines/config/pointnav/ppo_pointnav.yaml` config has better hyperparameters for large scale training and loads the [Gibson PointGoal Navigation Dataset](/README.md#datasets) instead of the test scenes.
 Change the `/benchmark/nav/pointnav: pointnav_gibson` in `habitat_baselines/config/pointnav/ppo_pointnav.yaml` to `/benchmark/nav/pointnav: pointnav_mp3d` in the defaults list for training on [MatterPort3D PointGoal Navigation Dataset](/README.md#datasets).
 
+### Hierarchical Reinforcement Learning (HRL)
+
+We provide a two-layer hierarchical policy class, consisting of a low-level skill that moves the robot, and a high-level policy that reasons about which low-level skill to use in the current state. This can be especially powerful in long-horizon mobile manipulation tasks, like those introduced in [Habitat2.0](https://arxiv.org/abs/2106.14405). Both the low- and high- level can be either learned or an oracle. For oracle high-level we use [PDDL](https://planning.wiki/guide/whatis/pddl), and for oracle low-level we use instantaneous transitions, with the environment set to the final desired state. Additionally, for navigation, we provide an oracle navigation skill that uses A-star and the map of the environment to move the robot to its goal.
+
+To run the following examples, you need the [ReplicaCAD dataset](https://github.com/facebookresearch/habitat-sim/blob/main/DATASETS.md#replicacad). 
+
+To train a high-level policy, while using pre-learned low-level skills (SRL baseline from [Habitat2.0](https://arxiv.org/abs/2106.14405)), you can run:
+
+```bash
+python -u habitat_baselines/run.py \
+  --exp-config habitat_baselines/config/rearrange/rl_hl_srl_onav.yaml \
+  --run-type train
+```
+To run a rearrangement episode with oracle in both low- and high-level, you can run:
+
+```bash
+python -u habitat_baselines/run.py \
+  --exp-config habitat_baselines/config/rearrange/rl_hl_srl_onav.yaml \
+  --run-type eval \
+  habitat_baselines/rl/policy=hierarchical_tp_noop_onav
+```
+
+To change the task (like set table) that you train your skills on, you can change the line `/habitat/task/rearrange: rearrange_easy` to `/habitat/task/rearrange: set_table` in the defaults of your config.
+
+
 ### Classic
 
 **SLAM based**

habitat-baselines/habitat_baselines/common/baseline_registry.py

@@ -136,5 +136,21 @@ def register_auxiliary_loss(
     def get_auxiliary_loss(cls, name: str):
         return cls._get_impl("aux_loss", name)
 
+    @classmethod
+    def register_storage(cls, to_register=None, *, name: Optional[str] = None):
+        return cls._register_impl("storage", to_register, name)
+
+    @classmethod
+    def get_storage(cls, name: str):
+        return cls._get_impl("storage", name)
+
+    @classmethod
+    def register_updater(cls, to_register=None, *, name: Optional[str] = None):
+        return cls._register_impl("updater", to_register, name)
+
+    @classmethod
+    def get_updater(cls, name: str):
+        return cls._get_impl("updater", name)
+
 
 baseline_registry = BaselineRegistry()

habitat-baselines/habitat_baselines/common/rollout_storage.py

@@ -10,13 +10,15 @@
 import numpy as np
 import torch
 
+from habitat_baselines.common.baseline_registry import baseline_registry
 from habitat_baselines.common.tensor_dict import DictTree, TensorDict
 from habitat_baselines.rl.models.rnn_state_encoder import (
     build_pack_info_from_dones,
     build_rnn_build_seq_info,
 )
 
 
+@baseline_registry.register_storage
 class RolloutStorage:
     r"""Class for storing rollout information for RL trainers."""
 
@@ -116,6 +118,7 @@ def insert(
         rewards=None,
         next_masks=None,
         buffer_index: int = 0,
+        **kwargs,
     ):
         if not self.is_double_buffered:
             assert buffer_index == 0

habitat-baselines/habitat_baselines/config/default_structured_configs.py

@@ -219,10 +219,31 @@ class Eq2CubeConfig(ObsTransformConfig):
 )
 
 
+@dataclass
+class HrlDefinedSkill(HabitatBaselinesBaseConfig):
+    skill_name: str = MISSING
+    name: str = "PointNavResNetPolicy"
+    action_distribution_type: str = "gaussian"
+    load_ckpt_file: str = ""
+    max_skill_steps: int = 200
+    force_end_on_timeout: bool = True
+    force_config_file: str = ""
+    at_resting_threshold: float = 0.15
+    apply_postconds: bool = False
+    obs_skill_inputs: List[str] = field(default_factory=list)
+    obs_skill_input_dim: int = 3
+    start_zone_radius: float = 0.3
+    # For the oracle navigation skill
+    nav_action_name: str = "base_velocity"
+    stop_thresh: float = 0.001
+    # For the reset_arm_skill
+    reset_joint_state: List[float] = MISSING
+
+
 @dataclass
 class HierarchicalPolicy(HabitatBaselinesBaseConfig):
     high_level_policy: Dict[str, Any] = MISSING
-    defined_skills: Dict[str, Any] = field(default_factory=dict)
+    defined_skills: Dict[str, HrlDefinedSkill] = field(default_factory=dict)
     use_skills: Dict[str, str] = field(default_factory=dict)
 
 
@@ -383,6 +404,8 @@ class HabitatBaselinesConfig(HabitatBaselinesBaseConfig):
     # )
     # cmd_trailing_opts: List[str] = field(default_factory=list)
     trainer_name: str = "ppo"
+    updater_name: str = "PPO"
+    distrib_updater_name: str = "DDPPO"
     torch_gpu_id: int = 0
     video_render_views: List[str] = field(default_factory=list)
     tensorboard_dir: str = "tb"
@@ -394,6 +417,7 @@ class HabitatBaselinesConfig(HabitatBaselinesBaseConfig):
     eval_ckpt_path_dir: str = "data/checkpoints"
     num_environments: int = 16
     num_processes: int = -1  # deprecated
+    rollout_storage: str = "RolloutStorage"
     checkpoint_folder: str = "data/checkpoints"
     num_updates: int = 10000
     num_checkpoints: int = 10

habitat-baselines/habitat_baselines/config/habitat_baselines/rl/policy/hierarchical_srl_onav.yaml

@@ -0,0 +1,77 @@
+name: "HierarchicalPolicy"
+obs_transforms:
+  add_virtual_keys:
+    virtual_keys:
+      "goal_to_agent_gps_compass": 2
+hierarchical_policy:
+  high_level_policy:
+    name: "NeuralHighLevelPolicy"
+    allow_other_place: False
+    hidden_dim: 512
+    use_rnn: True
+    rnn_type: 'LSTM'
+    backbone: resnet18
+    normalize_visual_inputs: False
+    num_rnn_layers: 2
+    policy_input_keys:
+      - "robot_head_depth"
+  defined_skills:
+    open_cab:
+      skill_name: "ArtObjSkillPolicy"
+      name: "PointNavResNetPolicy"
+      load_ckpt_file: "data/models/open_cab.pth"
+
+    open_fridge:
+      skill_name: "ArtObjSkillPolicy"
+      load_ckpt_file: "data/models/open_fridge.pth"
+
+    close_cab:
+      skill_name: "ArtObjSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/close_cab.pth"
+
+    close_fridge:
+      skill_name: "ArtObjSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/close_fridge.pth"
+
+    pick:
+      skill_name: "PickSkillPolicy"
+      name: "PointNavResNetPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/pick.pth"
+
+    place:
+      skill_name: "PlaceSkillPolicy"
+      name: "PointNavResNetPolicy"
+      obs_skill_inputs: ["obj_goal_sensor"]
+      load_ckpt_file: "data/models/place.pth"
+
+    wait_skill:
+      skill_name: "WaitSkillPolicy"
+      max_skill_steps: -1
+      force_end_on_timeout: False
+
+    oracle_nav:
+      skill_name: "OracleNavPolicy"
+      obs_skill_inputs: ["obj_start_sensor", "abs_obj_start_sensor", "obj_goal_sensor", "abs_obj_goal_sensor"]
+      max_skill_steps: 300
+
+    reset_arm_skill:
+      skill_name: "ResetArmSkill"
+      max_skill_steps: 50
+      reset_joint_state: [-4.50e-01, -1.08e00, 9.95e-02, 9.38e-01, -7.88e-04, 1.57e00, 4.62e-03]
+      force_end_on_timeout: False
+
+  use_skills:
+    # Uncomment if you are also using these skills
+    # open_cab: "open_cab"
+    # open_fridge: "open_fridge"
+    # close_cab: "open_cab"
+    # close_fridge: "open_fridge"
+    pick: "pick"
+    place: "place"
+    nav: "oracle_nav"
+    nav_to_receptacle: "oracle_nav"
+    wait: "wait_skill"
+    reset_arm: "reset_arm_skill"

habitat-baselines/habitat_baselines/config/habitat_baselines/rl/policy/hierarchical_tp_noop_onav.yaml

@@ -0,0 +1,72 @@
+
+name: "HierarchicalPolicy"
+obs_transforms:
+  add_virtual_keys:
+    virtual_keys:
+      "goal_to_agent_gps_compass": 2
+hierarchical_policy:
+  high_level_policy:
+    name: "FixedHighLevelPolicy"
+    add_arm_rest: True
+
+  defined_skills:
+    open_cab:
+      skill_name: "NoopSkillPolicy"
+      max_skill_steps: 1
+      apply_postconds: True
+
+    open_fridge:
+      skill_name: "NoopSkillPolicy"
+      max_skill_steps: 1
+      apply_postconds: True
+
+    close_cab:
+      skill_name: "NoopSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      max_skill_steps: 1
+
+    close_fridge:
+      skill_name: "NoopSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      max_skill_steps: 1
+      apply_postconds: True
+
+    pick:
+      skill_name: "NoopSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      max_skill_steps: 1
+      apply_postconds: True
+
+    place:
+      skill_name: "NoopSkillPolicy"
+      obs_skill_inputs: ["obj_goal_sensor"]
+      max_skill_steps: 1
+      apply_postconds: True
+
+    nav_to_obj:
+      skill_name: "OracleNavPolicy"
+      obs_skill_inputs: ["obj_start_sensor", "abs_obj_start_sensor", "obj_goal_sensor", "abs_obj_goal_sensor"]
+      max_skill_steps: 300
+
+    wait_skill:
+      skill_name: "WaitSkillPolicy"
+      max_skill_steps: -1
+      force_end_on_timeout: False
+
+    reset_arm_skill:
+      skill_name: "ResetArmSkill"
+      max_skill_steps: 50
+      reset_joint_state: [-4.50e-01, -1.07e00, 9.95e-02, 9.38e-01, -7.88e-04, 1.57e00, 4.62e-03]
+      force_end_on_timeout: False
+
+  use_skills:
+    open_cab: "open_cab"
+    open_fridge: "open_fridge"
+    close_cab: "close_cab"
+    close_fridge: "close_fridge"
+    pick: "pick"
+    place: "place"
+    nav: "nav_to_obj"
+    nav_to_receptacle: "nav_to_obj"
+    wait: "wait_skill"
+    reset_arm: "reset_arm_skill"

habitat-baselines/habitat_baselines/config/habitat_baselines/rl/policy/hierarchical_tp_srl.yaml

@@ -0,0 +1,71 @@
+name: "HierarchicalPolicy"
+obs_transforms:
+  add_virtual_keys:
+    virtual_keys:
+      "goal_to_agent_gps_compass": 2
+hierarchical_policy:
+  high_level_policy:
+    name: "FixedHighLevelPolicy"
+    add_arm_rest: True
+  defined_skills:
+    open_cab:
+      skill_name: "ArtObjSkillPolicy"
+      name: "PointNavResNetPolicy"
+      load_ckpt_file: "data/models/open_cab.pth"
+
+    open_fridge:
+      skill_name: "ArtObjSkillPolicy"
+      skill_name: "ArtObjSkillPolicy"
+      load_ckpt_file: "data/models/open_fridge.pth"
+
+    close_cab:
+      skill_name: "ArtObjSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/close_cab.pth"
+
+    close_fridge:
+      skill_name: "ArtObjSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/close_fridge.pth"
+
+    pick:
+      skill_name: "PickSkillPolicy"
+      obs_skill_inputs: ["obj_start_sensor"]
+      load_ckpt_file: "data/models/pick.pth"
+
+    place:
+      skill_name: "PlaceSkillPolicy"
+      obs_skill_inputs: ["obj_goal_sensor"]
+      load_ckpt_file: "data/models/place.pth"
+
+    wait_skill:
+      skill_name: "WaitSkillPolicy"
+      max_skill_steps: -1
+      force_end_on_timeout: False
+
+    nav_to_obj:
+      skill_name: "NavSkillPolicy"
+      obs_skill_inputs: ["goal_to_agent_gps_compass"]
+      obs_skill_input_dim: 2
+      load_ckpt_file: "data/models/nav.pth"
+      max_skill_steps: 300
+      force_end_on_timeout: False
+
+    reset_arm_skill:
+      skill_name: "ResetArmSkill"
+      max_skill_steps: 50
+      reset_joint_state: [-4.50e-01, -1.08e00, 9.95e-02, 9.38e-01, -7.88e-04, 1.57e00, 4.62e-03]
+      force_end_on_timeout: False
+
+  use_skills:
+    # Uncomment if you are also using these skills
+    # open_cab: "open_cab"
+    # open_fridge: "open_fridge"
+    # close_cab: "open_cab"
+    # close_fridge: "open_fridge"
+    pick: "pick"
+    place: "place"
+    nav: "nav_to_obj"
+    nav_to_receptacle: "nav_to_obj"
+    wait: "wait_skill"
+    reset_arm: "reset_arm_skill"

habitat-baselines/habitat_baselines/config/habitat_baselines/rl/policy/monolithic.yaml

@@ -0,0 +1,4 @@
+name: "PointNavResNetPolicy"
+action_distribution_type: "gaussian"
+action_dist:
+  use_log_std: True