Hot Load and Chunked Save of Demo Files

com.unity.ml-agents/CHANGELOG.md

@@ -12,7 +12,11 @@ and this project adheres to
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 ### Minor Changes
 #### com.unity.ml-agents / com.unity.ml-agents.extensions (C#)
+- The Demonstration Recorder will record demonstrations in chunks if `Num Steps To Record` is set to a
+value greater than 0. (#5455)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
+- During a training run with Behavior Cloning and/or GAIL, any demonstration files added to a demonstration
+directory will be loaded automatically. (#5455)
 ### Bug Fixes
 #### com.unity.ml-agents / com.unity.ml-agents.extensions (C#)
 #### ml-agents / ml-agents-envs / gym-unity (Python)

com.unity.ml-agents/Runtime/Demonstrations/DemonstrationRecorder.cs

@@ -65,6 +65,8 @@ public class DemonstrationRecorder : MonoBehaviour
         const string k_DefaultDirectoryName = "Demonstrations";
         IFileSystem m_FileSystem;
 
+        private string currentfilePath;
+
         Agent m_Agent;
 
         void OnEnable()
@@ -81,8 +83,14 @@ void Update()
 
             LazyInitialize();
 
-            // Quit when num steps to record is reached
+            // Write new file when num steps is reached
             if (NumStepsToRecord > 0 && m_DemoWriter.NumSteps >= NumStepsToRecord)
+            {
+                Close();
+            }
+
+            // Quit when num steps to record is reached
+            if (false && NumStepsToRecord > 0 && m_DemoWriter.NumSteps >= NumStepsToRecord)
             {
                 Application.Quit(0);
 #if UNITY_EDITOR
@@ -120,7 +128,12 @@ internal DemonstrationWriter LazyInitialize(IFileSystem fileSystem = null)
 
             DemonstrationName = SanitizeName(DemonstrationName, MaxNameLength);
             var filePath = MakeDemonstrationFilePath(m_FileSystem, DemonstrationDirectory, DemonstrationName);
-            var stream = m_FileSystem.File.Create(filePath);
+            currentfilePath = filePath + ".tmp";
+            if (m_FileSystem.File.Exists(filePath))
+            {
+                m_FileSystem.File.Delete(filePath);
+            }
+            var stream = m_FileSystem.File.Create(currentfilePath);
             m_DemoWriter = new DemonstrationWriter(stream);
 
             AddDemonstrationWriterToAgent(m_DemoWriter);
@@ -188,6 +201,12 @@ public void Close()
 
                 m_DemoWriter.Close();
                 m_DemoWriter = null;
+                int fileExtPos = currentfilePath.LastIndexOf(".tmp");
+                if (fileExtPos >= 0)
+                {
+                    var newFilePath = currentfilePath.Substring(0, fileExtPos);
+                    m_FileSystem.File.Move(currentfilePath, newFilePath);
+                }
             }
         }
 

ml-agents/mlagents/trainers/demo_loader.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Tuple
+from typing import List, Set, Tuple
 import numpy as np
 from mlagents.trainers.buffer import AgentBuffer, BufferKey
 from mlagents_envs.communicator_objects.agent_info_action_pair_pb2 import (
@@ -13,88 +13,95 @@
     DemonstrationMetaProto,
 )
 from mlagents_envs.timers import timed, hierarchical_timer
+from mlagents_envs import logging_util
 from google.protobuf.internal.decoder import _DecodeVarint32  # type: ignore
 from google.protobuf.internal.encoder import _EncodeVarint  # type: ignore
 
 
 INITIAL_POS = 33
 SUPPORTED_DEMONSTRATION_VERSIONS = frozenset([0, 1])
 
+logger = logging_util.get_logger(__name__)
+
 
 @timed
 def make_demo_buffer(
-    pair_infos: List[AgentInfoActionPairProto],
+    all_pair_infos: List[List[AgentInfoActionPairProto]],
     behavior_spec: BehaviorSpec,
     sequence_length: int,
 ) -> AgentBuffer:
     # Create and populate buffer using experiences
     demo_raw_buffer = AgentBuffer()
     demo_processed_buffer = AgentBuffer()
-    for idx, current_pair_info in enumerate(pair_infos):
-        if idx > len(pair_infos) - 2:
-            break
-        next_pair_info = pair_infos[idx + 1]
-        current_decision_step, current_terminal_step = steps_from_proto(
-            [current_pair_info.agent_info], behavior_spec
-        )
-        next_decision_step, next_terminal_step = steps_from_proto(
-            [next_pair_info.agent_info], behavior_spec
-        )
-        previous_action = (
-            np.array(
-                pair_infos[idx].action_info.vector_actions_deprecated, dtype=np.float32
+    for pair_infos in all_pair_infos:
+        for idx, current_pair_info in enumerate(pair_infos):
+            if idx > len(pair_infos) - 2:
+                break
+            next_pair_info = pair_infos[idx + 1]
+            current_decision_step, current_terminal_step = steps_from_proto(
+                [current_pair_info.agent_info], behavior_spec
             )
-            * 0
-        )
-        if idx > 0:
-            previous_action = np.array(
-                pair_infos[idx - 1].action_info.vector_actions_deprecated,
-                dtype=np.float32,
+            next_decision_step, next_terminal_step = steps_from_proto(
+                [next_pair_info.agent_info], behavior_spec
             )
-
-        next_done = len(next_terminal_step) == 1
-        next_reward = 0
-        if len(next_terminal_step) == 1:
-            next_reward = next_terminal_step.reward[0]
-        else:
-            next_reward = next_decision_step.reward[0]
-        current_obs = None
-        if len(current_terminal_step) == 1:
-            current_obs = list(current_terminal_step.values())[0].obs
-        else:
-            current_obs = list(current_decision_step.values())[0].obs
-
-        demo_raw_buffer[BufferKey.DONE].append(next_done)
-        demo_raw_buffer[BufferKey.ENVIRONMENT_REWARDS].append(next_reward)
-        for i, obs in enumerate(current_obs):
-            demo_raw_buffer[ObsUtil.get_name_at(i)].append(obs)
-        if (
-            len(current_pair_info.action_info.continuous_actions) == 0
-            and len(current_pair_info.action_info.discrete_actions) == 0
-        ):
-            if behavior_spec.action_spec.continuous_size > 0:
-                demo_raw_buffer[BufferKey.CONTINUOUS_ACTION].append(
-                    current_pair_info.action_info.vector_actions_deprecated
+            previous_action = (
+                np.array(
+                    pair_infos[idx].action_info.vector_actions_deprecated,
+                    dtype=np.float32,
                 )
-            else:
-                demo_raw_buffer[BufferKey.DISCRETE_ACTION].append(
-                    current_pair_info.action_info.vector_actions_deprecated
-                )
-        else:
-            if behavior_spec.action_spec.continuous_size > 0:
-                demo_raw_buffer[BufferKey.CONTINUOUS_ACTION].append(
-                    current_pair_info.action_info.continuous_actions
+                * 0
+            )
+            if idx > 0:
+                previous_action = np.array(
+                    pair_infos[idx - 1].action_info.vector_actions_deprecated,
+                    dtype=np.float32,
                 )
-            if behavior_spec.action_spec.discrete_size > 0:
-                demo_raw_buffer[BufferKey.DISCRETE_ACTION].append(
-                    current_pair_info.action_info.discrete_actions
+
+            next_done = len(next_terminal_step) == 1
+            next_reward = 0
+            if len(next_terminal_step) == 1:
+                next_reward = next_terminal_step.reward[0]
+            else:
+                next_reward = next_decision_step.reward[0]
+            current_obs = None
+            if len(current_terminal_step) == 1:
+                current_obs = list(current_terminal_step.values())[0].obs
+            else:
+                current_obs = list(current_decision_step.values())[0].obs
+
+            demo_raw_buffer[BufferKey.DONE].append(next_done)
+            demo_raw_buffer[BufferKey.ENVIRONMENT_REWARDS].append(next_reward)
+            for i, obs in enumerate(current_obs):
+                demo_raw_buffer[ObsUtil.get_name_at(i)].append(obs)
+            if (
+                len(current_pair_info.action_info.continuous_actions) == 0
+                and len(current_pair_info.action_info.discrete_actions) == 0
+            ):
+                if behavior_spec.action_spec.continuous_size > 0:
+                    demo_raw_buffer[BufferKey.CONTINUOUS_ACTION].append(
+                        current_pair_info.action_info.vector_actions_deprecated
+                    )
+                else:
+                    demo_raw_buffer[BufferKey.DISCRETE_ACTION].append(
+                        current_pair_info.action_info.vector_actions_deprecated
+                    )
+            else:
+                if behavior_spec.action_spec.continuous_size > 0:
+                    demo_raw_buffer[BufferKey.CONTINUOUS_ACTION].append(
+                        current_pair_info.action_info.continuous_actions
+                    )
+                if behavior_spec.action_spec.discrete_size > 0:
+                    demo_raw_buffer[BufferKey.DISCRETE_ACTION].append(
+                        current_pair_info.action_info.discrete_actions
+                    )
+            demo_raw_buffer[BufferKey.PREV_ACTION].append(previous_action)
+            if next_done:
+                demo_raw_buffer.resequence_and_append(
+                    demo_processed_buffer,
+                    batch_size=None,
+                    training_length=sequence_length,
                 )
-        demo_raw_buffer[BufferKey.PREV_ACTION].append(previous_action)
-        if next_done:
-            demo_raw_buffer.resequence_and_append(
-                demo_processed_buffer, batch_size=None, training_length=sequence_length
-            )
-            demo_raw_buffer.reset_agent()
+                demo_raw_buffer.reset_agent()
     demo_raw_buffer.resequence_and_append(
         demo_processed_buffer, batch_size=None, training_length=sequence_length
     )
@@ -104,15 +111,17 @@ def make_demo_buffer(
 @timed
 def demo_to_buffer(
     file_path: str, sequence_length: int, expected_behavior_spec: BehaviorSpec = None
-) -> Tuple[BehaviorSpec, AgentBuffer]:
+) -> Tuple[BehaviorSpec, AgentBuffer, List[str]]:
     """
     Loads demonstration file and uses it to fill training buffer.
     :param file_path: Location of demonstration file (.demo).
     :param sequence_length: Length of trajectories to fill buffer.
     :return:
     """
-    behavior_spec, info_action_pair, _ = load_demonstration(file_path)
-    demo_buffer = make_demo_buffer(info_action_pair, behavior_spec, sequence_length)
+    behavior_spec, info_action_pairs, file_paths = load_demonstration(file_path)
+    if len(file_paths) == 0:
+        return None, AgentBuffer(), file_paths
+    demo_buffer = make_demo_buffer(info_action_pairs, behavior_spec, sequence_length)
     if expected_behavior_spec:
         # check action dimensions in demonstration match
         if behavior_spec.action_spec != expected_behavior_spec.action_spec:
@@ -140,7 +149,7 @@ def demo_to_buffer(
                         f"The shape {demo_obs} for observation {i} in demonstration \
                         do not match the policy's {policy_obs}."
                     )
-    return behavior_spec, demo_buffer
+    return behavior_spec, demo_buffer, file_paths
 
 
 def get_demo_files(path: str) -> List[str]:
@@ -162,7 +171,7 @@ def get_demo_files(path: str) -> List[str]:
             if name.endswith(".demo")
         ]
         if not paths:
-            raise ValueError("There are no '.demo' files in the provided directory.")
+            logger.debug("There are no '.demo' files in the provided directory.")
         return paths
     else:
         raise FileNotFoundError(
@@ -172,22 +181,27 @@ def get_demo_files(path: str) -> List[str]:
 
 @timed
 def load_demonstration(
-    file_path: str,
-) -> Tuple[BehaviorSpec, List[AgentInfoActionPairProto], int]:
+    file_path: str, exclusions: set = None
+) -> Tuple[BehaviorSpec, List[AgentInfoActionPairProto], List[str]]:
     """
     Loads and parses a demonstration file.
     :param file_path: Location of demonstration file (.demo).
     :return: BrainParameter and list of AgentInfoActionPairProto containing demonstration data.
     """
 
     # First 32 bytes of file dedicated to meta-data.
-    file_paths = get_demo_files(file_path)
+    file_paths = set(get_demo_files(file_path))
+    if exclusions is not None:
+        file_paths = file_paths - exclusions
     behavior_spec = None
     brain_param_proto = None
-    info_action_pairs = []
     total_expected = 0
+    # We divide it up in files so that we don't accidentally merge them
+    info_action_pairs: List[List[AgentInfoActionPairProto]] = []
     for _file_path in file_paths:
         with open(_file_path, "rb") as fp:
+            expected_in_file = 0
+            _info_pairs_per_file = []
             with hierarchical_timer("read_file"):
                 data = fp.read()
             next_pos, pos, obs_decoded = 0, 0, 0
@@ -201,9 +215,10 @@ def load_demonstration(
                         not in SUPPORTED_DEMONSTRATION_VERSIONS
                     ):
                         raise RuntimeError(
-                            f"Can't load Demonstration data from an unsupported version ({meta_data_proto.api_version})"
+                            f"Can't load Demonstration data from an unsupported version \
+                                ({meta_data_proto.api_version})"
                         )
-                    total_expected += meta_data_proto.number_steps
+                    expected_in_file += meta_data_proto.number_steps
                     pos = INITIAL_POS
                 if obs_decoded == 1:
                     brain_param_proto = BrainParametersProto()
@@ -216,16 +231,18 @@ def load_demonstration(
                         behavior_spec = behavior_spec_from_proto(
                             brain_param_proto, agent_info_action.agent_info
                         )
-                    info_action_pairs.append(agent_info_action)
-                    if len(info_action_pairs) == total_expected:
+                    _info_pairs_per_file.append(agent_info_action)
+                    if len(_info_pairs_per_file) == expected_in_file:
                         break
                     pos += next_pos
                 obs_decoded += 1
-    if not behavior_spec:
+        info_action_pairs.append(_info_pairs_per_file)
+        total_expected += expected_in_file
+    if not behavior_spec and total_expected > 0:
         raise RuntimeError(
             f"No BrainParameters found in demonstration file at {file_path}."
         )
-    return behavior_spec, info_action_pairs, total_expected
+    return behavior_spec, info_action_pairs, list(file_paths)
 
 
 def write_delimited(f, message):
@@ -244,3 +261,47 @@ def write_demo(demo_path, meta_data_proto, brain_param_proto, agent_info_protos)
 
         for agent in agent_info_protos:
             write_delimited(f, agent)
+
+
+class DemoManager:
+    def __init__(
+        self,
+        path: str,
+        sequence_length: int = 1,
+        expected_bspec: BehaviorSpec = None,
+        demo_buffer_size: int = 1000000,
+    ) -> None:
+        self.path = path
+        self._seq_len = sequence_length
+        self._loaded_files: Set[str] = set()
+        _, self._demo_buffer, file_paths = demo_to_buffer(
+            path, sequence_length, expected_bspec
+        )
+        self._loaded_files.update(file_paths)
+        self._max_demo_buffer_size = demo_buffer_size
+        if len(file_paths) == 0:
+            self._demo_buffer = AgentBuffer()
+            logger.warn(f"No demos found in {path}. Continuing to look for new files.")
+
+    @property
+    def demo_buffer(self) -> AgentBuffer:
+        return self._demo_buffer
+
+    def refresh(self) -> int:
+        bspec, loaded_demos, loaded_paths = load_demonstration(
+            self.path, self._loaded_files
+        )
+        if loaded_paths:
+            new_demos = make_demo_buffer(loaded_demos, bspec, self._seq_len)
+            new_demos.resequence_and_append(
+                self._demo_buffer, training_length=self._seq_len
+            )
+            self._loaded_files.update(loaded_paths)
+            num_new_exp = new_demos.num_experiences
+            if self._demo_buffer.num_experiences > self._max_demo_buffer_size:
+                self._demo_buffer.truncate(
+                    int(self._max_demo_buffer_size * 0.8), self._seq_len
+                )
+            return num_new_exp
+        else:
+            return 0