Comparing for eventual rebase

.circleci/config.yml

@@ -203,12 +203,15 @@ jobs:
       - run:
           name: Download test data
           command: |
+              # Disable clone protection for git lfs
+              export GIT_CLONE_PROTECTION_ACTIVE=false
+
               sudo apt install git-lfs
               export PATH=$HOME/miniconda/bin:/usr/local/cuda/bin:$PATH
               . activate habitat
               git lfs install
               conda install -y gitpython git-lfs
-              python -m habitat_sim.utils.datasets_download --uids ci_test_assets franka_panda hab_spot_arm hab3_bench_assets --data-path habitat-sim/data/ --no-replace --no-prune
+              python -m habitat_sim.utils.datasets_download --uids ci_test_assets franka_panda hab_spot_arm hab3_bench_assets ycb rearrange_dataset_v2 --data-path habitat-sim/data/ --no-replace --no-prune
       - run:
           name: Run sim benchmark
           command: |
@@ -281,6 +284,7 @@ jobs:
               . activate habitat; cd habitat-lab
               export PYTHONPATH=.:$PYTHONPATH
               export MULTI_PROC_OFFSET=0 && export MAGNUM_LOG=quiet && export HABITAT_SIM_LOG=quiet
+              export GIT_CLONE_PROTECTION_ACTIVE=false
               python -m habitat_sim.utils.datasets_download --uids hab3-episodes hab3_bench_assets habitat_humanoids hab_spot_arm ycb --data-path data/ --no-replace --no-prune
               python -m pytest habitat-hitl/test
       - run:

.gitignore

@@ -100,3 +100,6 @@ data
 /sandbox
 /plots/outputs
 /wandb/
+
+# Autogenerated navmesh files
+*.navmesh

examples/hitl/rearrange_v2/app_data.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict
+
+from habitat_hitl.core.types import ConnectionRecord
+
+
+class AppData:
+    """
+    RearrangeV2 application data shared by all states.
+    """
+
+    def __init__(self, max_user_count: int):
+        self.max_user_count = max_user_count
+        self.connected_users: Dict[int, ConnectionRecord] = {}

examples/hitl/rearrange_v2/app_state_base.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from typing import Optional
+
+from app_data import AppData
+
+from habitat_hitl.app_states.app_service import AppService
+from habitat_hitl.app_states.app_state_abc import AppState
+from habitat_hitl.core.text_drawer import TextOnScreenAlignment
+from habitat_hitl.core.user_mask import Mask
+
+
+class AppStateBase(AppState):
+    def __init__(
+        self,
+        app_service: AppService,
+        app_data: AppData,
+    ):
+        self._app_service = app_service
+        self._app_data = app_data
+        self._cancel = False
+        self._time_since_last_connection = 0
+        self._save_keyframes = True
+
+    def on_enter(self):
+        print(f"Entering state: {type(self)}")
+
+    def on_exit(self):
+        print(f"Exiting state: {type(self)}")
+
+    def try_cancel(self):
+        self._cancel = True
+
+    def get_next_state(self) -> Optional[AppStateBase]:
+        pass
+
+    def on_environment_reset(self, episode_recorder_dict):
+        pass
+
+    def sim_update(self, dt: float, post_sim_update_dict):
+        pass
+
+    def record_state(self):
+        pass
+
+    def _status_message(self, message: str) -> None:
+        """Send a message to all users."""
+        if len(message) > 0:
+            self._app_service.text_drawer.add_text(
+                message,
+                TextOnScreenAlignment.TOP_CENTER,
+                text_delta_x=-280,
+                text_delta_y=-50,
+                destination_mask=Mask.ALL,
+            )
+
+    def _kick_all_users(self) -> None:
+        "Kick all users."
+        self._app_service.remote_client_state.kick(Mask.ALL)
+
+    def _is_server_gui_enabled(self) -> bool:
+        "Returns true if the local server GUI is available."
+        return (
+            not self._app_service.hitl_config.experimental.headless.do_headless
+        )

examples/hitl/rearrange_v2/app_state_end_session.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import shutil
+from typing import Optional
+
+from app_data import AppData
+from app_state_base import AppStateBase
+from app_states import create_app_state_reset
+from s3_upload import (
+    generate_unique_session_id,
+    make_s3_filename,
+    upload_file_to_s3,
+)
+from session import Session
+from util import get_top_down_view
+
+from habitat_hitl.app_states.app_service import AppService
+from habitat_hitl.core.serialize_utils import save_as_json_gzip
+from habitat_hitl.core.user_mask import Mask
+
+# Duration of the end session message, before users are kicked.
+SESSION_END_DELAY = 5.0
+
+
+class AppStateEndSession(AppStateBase):
+    """
+    * Indicate users that the session is terminated.
+    * Upload collected data.
+    """
+
+    def __init__(
+        self, app_service: AppService, app_data: AppData, session: Session
+    ):
+        super().__init__(app_service, app_data)
+        self._session = session
+        self._elapsed_time = 0.0
+        self._save_keyframes = False
+
+        self._status = "Session ended."
+        if len(session.error) > 0:
+            self._status += f"\nError: {session.error}"
+
+    def get_next_state(self) -> Optional[AppStateBase]:
+        if self._elapsed_time > SESSION_END_DELAY:
+            self._end_session()
+            return create_app_state_reset(self._app_service, self._app_data)
+        return None
+
+    def sim_update(self, dt: float, post_sim_update_dict):
+        # Top-down view.
+        cam_matrix = get_top_down_view(self._app_service.sim)
+        post_sim_update_dict["cam_transform"] = cam_matrix
+        self._app_service._client_message_manager.update_camera_transform(
+            cam_matrix, destination_mask=Mask.ALL
+        )
+
+        self._status_message(self._status)
+        self._elapsed_time += dt
+
+    def _end_session(self):
+        session = self._session
+        if session is None:
+            print("Null session. Skipping S3 upload.")
+            return
+
+        # Finalize session.
+        if self._session.error == "":
+            session.success = True
+        session.session_recorder.end_session(self._session.error)
+
+        # Get data collection parameters.
+        try:
+            config = self._app_service.config
+            data_collection_config = config.rearrange_v2.data_collection
+            s3_path = data_collection_config.s3_path
+            s3_subdir = "complete" if session.success else "incomplete"
+            s3_path = os.path.join(s3_path, s3_subdir)
+
+            # Use the port as a discriminator for when there are multiple concurrent servers.
+            output_folder_suffix = str(config.habitat_hitl.networking.port)
+            output_folder = f"data/output_{output_folder_suffix}"
+
+            output_file_name = data_collection_config.output_file_name
+            output_file = f"{output_file_name}.json.gz"
+
+        except Exception as e:
+            print(f"Invalid data collection config. Skipping S3 upload. {e}")
+            return
+
+        # Delete previous output directory
+        if os.path.exists(output_folder):
+            shutil.rmtree(output_folder)
+
+        # Create new output directory
+        os.makedirs(output_folder)
+        json_path = os.path.join(output_folder, output_file)
+        save_as_json_gzip(session.session_recorder, json_path)
+
+        # Generate unique session ID
+        session_id = generate_unique_session_id(
+            session.episode_ids, session.connection_records
+        )
+
+        # Upload output directory
+        orig_file_names = [
+            f
+            for f in os.listdir(output_folder)
+            if os.path.isfile(os.path.join(output_folder, f))
+        ]
+        for orig_file_name in orig_file_names:
+            local_file_path = os.path.join(output_folder, orig_file_name)
+            s3_file_name = make_s3_filename(session_id, orig_file_name)
+            upload_file_to_s3(local_file_path, s3_file_name, s3_path)

examples/hitl/rearrange_v2/app_state_load_episode.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+from app_data import AppData
+from app_state_base import AppStateBase
+from app_states import (
+    create_app_state_cancel_session,
+    create_app_state_end_session,
+    create_app_state_start_screen,
+)
+from session import Session
+from util import get_top_down_view
+
+from habitat_hitl.app_states.app_service import AppService
+from habitat_hitl.core.user_mask import Mask
+
+
+class AppStateLoadEpisode(AppStateBase):
+    """
+    Load an episode.
+    A loading screen is shown while the content loads.
+    * If a next episode exists, launch RearrangeV2.
+    * If all episodes are done, end session.
+    * If any user disconnects, cancel the session.
+    """
+
+    def __init__(
+        self, app_service: AppService, app_data: AppData, session: Session
+    ):
+        super().__init__(app_service, app_data)
+        self._session = session
+        self._loading = True
+        self._session_ended = False
+        self._frame_number = 0
+        self._save_keyframes = False
+
+    def get_next_state(self) -> Optional[AppStateBase]:
+        if self._cancel:
+            return create_app_state_cancel_session(
+                self._app_service,
+                self._app_data,
+                self._session,
+                "User disconnected.",
+            )
+        if self._session_ended:
+            return create_app_state_end_session(
+                self._app_service, self._app_data, self._session
+            )
+        # When all clients finish loading, show the start screen.
+        if not self._loading:
+            return create_app_state_start_screen(
+                self._app_service, self._app_data, self._session
+            )
+        return None
+
+    def sim_update(self, dt: float, post_sim_update_dict):
+        self._status_message("Loading...")
+
+        # Skip a frame so that the status message reaches the client before the server loads the scene and blocks.
+        if self._frame_number == 1:
+            self._increment_episode()
+        # Once the scene loaded, show a top-down view.
+        elif self._frame_number > 1:
+            cam_matrix = get_top_down_view(self._app_service.sim)
+            post_sim_update_dict["cam_transform"] = cam_matrix
+            self._app_service._client_message_manager.update_camera_transform(
+                cam_matrix, destination_mask=Mask.ALL
+            )
+            # Wait for clients to signal that content finished loading on their end.
+            # HACK: The server isn't immediately aware that clients are loading. For now, we simply skip some frames.
+            # TODO: Use the keyframe ID from 'ClientMessageManager.set_server_keyframe_id()' to find the when the loading state is up-to-date.
+            if self._frame_number > 20:
+                any_client_loading = False
+                for user_index in range(self._app_data.max_user_count):
+                    if self._app_service.remote_client_state._client_loading[
+                        user_index
+                    ]:
+                        any_client_loading = True
+                        break
+                if not any_client_loading:
+                    self._loading = False
+
+        self._frame_number += 1
+
+    def _increment_episode(self):
+        session = self._session
+        assert session.episode_ids is not None
+        if session.current_episode_index < len(session.episode_ids):
+            self._set_episode(session.current_episode_index)
+            session.current_episode_index += 1
+        else:
+            self._session_ended = True
+
+    def _set_episode(self, episode_index: int):
+        session = self._session
+
+        # Set the ID of the next episode to play in lab.
+        next_episode_id = session.episode_ids[episode_index]
+        print(f"Next episode index: {next_episode_id}.")
+        try:
+            next_episode_index = int(next_episode_id)
+            self._app_service.episode_helper.set_next_episode_by_index(
+                next_episode_index
+            )
+        except Exception as e:
+            print(f"ERROR: Invalid episode index {next_episode_id}. {e}")
+            print("Loading episode index 0.")
+            self._app_service.episode_helper.set_next_episode_by_index(0)
+
+        # Once an episode ID has been set, lab needs to be reset to load the episode.
+        self._app_service.end_episode(do_reset=True)
+
+        # Signal the clients that the scene has changed.
+        client_message_manager = self._app_service.client_message_manager
+        if client_message_manager:
+            client_message_manager.signal_scene_change(Mask.ALL)
+
+        # Save a keyframe. This propagates the new content to the clients, initiating client-side loading.
+        # Beware that the client "loading" state won't immediately be visible to the server.
+        self._app_service.sim.gfx_replay_manager.save_keyframe()