diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md index ccab15150d9..a7b64ab8755 100755 --- a/com.unity.ml-agents/CHANGELOG.md +++ b/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) diff --git a/com.unity.ml-agents/Runtime/Demonstrations/DemonstrationRecorder.cs b/com.unity.ml-agents/Runtime/Demonstrations/DemonstrationRecorder.cs index 87d02b3fcaa..52eb2dc6104 100644 --- a/com.unity.ml-agents/Runtime/Demonstrations/DemonstrationRecorder.cs +++ b/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); + } } } diff --git a/ml-agents/mlagents/trainers/demo_loader.py b/ml-agents/mlagents/trainers/demo_loader.py index eba0200b883..c8c119e7c92 100644 --- a/ml-agents/mlagents/trainers/demo_loader.py +++ b/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,6 +13,7 @@ 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 @@ -20,81 +21,87 @@ 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,8 +181,8 @@ 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). @@ -181,13 +190,18 @@ def load_demonstration( """ # 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 diff --git a/ml-agents/mlagents/trainers/tests/test_demo_loader.py b/ml-agents/mlagents/trainers/tests/test_demo_loader.py index 7b2bb549c34..5020849b46d 100644 --- a/ml-agents/mlagents/trainers/tests/test_demo_loader.py +++ b/ml-agents/mlagents/trainers/tests/test_demo_loader.py @@ -1,5 +1,6 @@ import io import os +import shutil from unittest import mock import numpy as np import pytest @@ -13,6 +14,7 @@ setup_test_behavior_specs, ) from mlagents.trainers.demo_loader import ( + DemoManager, load_demonstration, demo_to_buffer, get_demo_files, @@ -26,13 +28,12 @@ def test_load_demo(): path_prefix = os.path.dirname(os.path.abspath(__file__)) - behavior_spec, pair_infos, total_expected = load_demonstration( - path_prefix + "/test.demo" - ) + behavior_spec, pair_infos, _ = load_demonstration(path_prefix + "/test.demo") assert np.sum(behavior_spec.observation_specs[0].shape) == 8 - assert len(pair_infos) == total_expected + total_expected = 2144 + assert len(pair_infos[0]) == total_expected - _, demo_buffer = demo_to_buffer(path_prefix + "/test.demo", 1, BEHAVIOR_SPEC) + _, demo_buffer, _ = demo_to_buffer(path_prefix + "/test.demo", 1, BEHAVIOR_SPEC) assert ( len(demo_buffer[BufferKey.CONTINUOUS_ACTION]) == total_expected - 1 or len(demo_buffer[BufferKey.DISCRETE_ACTION]) == total_expected - 1 @@ -41,16 +42,21 @@ def test_load_demo(): def test_load_demo_dir(): path_prefix = os.path.dirname(os.path.abspath(__file__)) - behavior_spec, pair_infos, total_expected = load_demonstration( - path_prefix + "/test_demo_dir" - ) + behavior_spec, pair_infos, _ = load_demonstration(path_prefix + "/test_demo_dir") assert np.sum(behavior_spec.observation_specs[0].shape) == 8 - assert len(pair_infos) == total_expected - - _, demo_buffer = demo_to_buffer(path_prefix + "/test_demo_dir", 1, BEHAVIOR_SPEC) - assert ( - len(demo_buffer[BufferKey.CONTINUOUS_ACTION]) == total_expected - 1 - or len(demo_buffer[BufferKey.DISCRETE_ACTION]) == total_expected - 1 + total_expected_per_file = 25 + for pair_infos_per_file in pair_infos: + assert len(pair_infos_per_file) == total_expected_per_file + assert len(pair_infos) == 3 + + _, demo_buffer, _ = demo_to_buffer(path_prefix + "/test_demo_dir", 1, BEHAVIOR_SPEC) + # Each file should produce num_infos - 1 samples in the buffer. + assert len(demo_buffer[BufferKey.CONTINUOUS_ACTION]) == ( + total_expected_per_file - 1 + ) * len(pair_infos) or len(demo_buffer[BufferKey.DISCRETE_ACTION]) == ( + total_expected_per_file - 1 + ) * len( + pair_infos ) @@ -94,18 +100,12 @@ def test_edge_cases(): with pytest.raises(FileNotFoundError): get_demo_files(os.path.join(path_prefix, "nonexistent_directory")) with tempfile.TemporaryDirectory() as tmpdirname: - # empty directory - with pytest.raises(ValueError): - get_demo_files(tmpdirname) # invalid file invalid_fname = os.path.join(tmpdirname, "mydemo.notademo") with open(invalid_fname, "w") as f: f.write("I'm not a demo") with pytest.raises(ValueError): get_demo_files(invalid_fname) - # invalid directory - with pytest.raises(ValueError): - get_demo_files(tmpdirname) # valid file valid_fname = os.path.join(tmpdirname, "mydemo.demo") with open(valid_fname, "w") as f: @@ -130,3 +130,29 @@ def test_unsupported_version_raises_error(mock_get_demo_files): with mock.patch("builtins.open", m): with pytest.raises(RuntimeError): load_demonstration("foo") + + +def test_demo_manager(): + path_prefix = os.path.dirname(os.path.abspath(__file__)) + with tempfile.TemporaryDirectory() as tmpdirname: + # Initialize with empty dir + demo_man = DemoManager(tmpdirname) + assert demo_man.demo_buffer.num_experiences == 0 + # Add a demo file + shutil.copy2( + path_prefix + "/test_demo_dir/test4.demo", + os.path.join(tmpdirname, "test4.demo"), + ) + demo_man.refresh() + assert demo_man.demo_buffer.num_experiences == 24 + # Add another demo file + shutil.copy2( + path_prefix + "/test_demo_dir/test5.demo", + os.path.join(tmpdirname, "test5.demo"), + ) + demo_man.refresh() + assert demo_man.demo_buffer.num_experiences == 48 + + # Try initializing a new demo manager with a full dir + demo_man = DemoManager(tmpdirname) + assert demo_man.demo_buffer.num_experiences == 48 diff --git a/ml-agents/mlagents/trainers/tests/torch/test_reward_providers/test_gail.py b/ml-agents/mlagents/trainers/tests/torch/test_reward_providers/test_gail.py index 974faa33262..f21378c4e78 100644 --- a/ml-agents/mlagents/trainers/tests/torch/test_reward_providers/test_gail.py +++ b/ml-agents/mlagents/trainers/tests/torch/test_reward_providers/test_gail.py @@ -1,6 +1,7 @@ from typing import Any import numpy as np import pytest +from unittest import mock from unittest.mock import patch from mlagents.torch_utils import torch import os @@ -74,16 +75,19 @@ def test_factory(behavior_spec: BehaviorSpec) -> None: ) @pytest.mark.parametrize("use_actions", [False, True]) @patch( - "mlagents.trainers.torch.components.reward_providers.gail_reward_provider.demo_to_buffer" + "mlagents.trainers.torch.components.reward_providers.gail_reward_provider.DemoManager" ) def test_reward_decreases( - demo_to_buffer: Any, use_actions: bool, behavior_spec: BehaviorSpec, seed: int + mock_demo_manager_fn: Any, use_actions: bool, behavior_spec: BehaviorSpec, seed: int ) -> None: np.random.seed(seed) torch.manual_seed(seed) buffer_expert = create_agent_buffer(behavior_spec, 1000) buffer_policy = create_agent_buffer(behavior_spec, 1000) - demo_to_buffer.return_value = None, buffer_expert + mock_demo_manager = mock.Mock() + mock_demo_manager.demo_buffer = buffer_expert + mock_demo_manager.refresh.return_value = 0 + mock_demo_manager_fn.return_value = mock_demo_manager gail_settings = GAILSettings( demo_path="", learning_rate=0.005, use_vail=False, use_actions=use_actions ) @@ -129,16 +133,19 @@ def test_reward_decreases( ) @pytest.mark.parametrize("use_actions", [False, True]) @patch( - "mlagents.trainers.torch.components.reward_providers.gail_reward_provider.demo_to_buffer" + "mlagents.trainers.torch.components.reward_providers.gail_reward_provider.DemoManager" ) def test_reward_decreases_vail( - demo_to_buffer: Any, use_actions: bool, behavior_spec: BehaviorSpec, seed: int + mock_demo_manager_fn: Any, use_actions: bool, behavior_spec: BehaviorSpec, seed: int ) -> None: np.random.seed(seed) torch.manual_seed(seed) buffer_expert = create_agent_buffer(behavior_spec, 1000) buffer_policy = create_agent_buffer(behavior_spec, 1000) - demo_to_buffer.return_value = None, buffer_expert + mock_demo_manager = mock.Mock() + mock_demo_manager.demo_buffer = buffer_expert + mock_demo_manager.refresh.return_value = 0 + mock_demo_manager_fn.return_value = mock_demo_manager gail_settings = GAILSettings( demo_path="", learning_rate=0.005, use_vail=True, use_actions=use_actions ) diff --git a/ml-agents/mlagents/trainers/torch/components/bc/module.py b/ml-agents/mlagents/trainers/torch/components/bc/module.py index b8879569b4c..3f02c598e59 100644 --- a/ml-agents/mlagents/trainers/torch/components/bc/module.py +++ b/ml-agents/mlagents/trainers/torch/components/bc/module.py @@ -3,13 +3,16 @@ from mlagents.torch_utils import torch from mlagents.trainers.policy.torch_policy import TorchPolicy -from mlagents.trainers.demo_loader import demo_to_buffer +from mlagents.trainers.demo_loader import DemoManager from mlagents.trainers.settings import BehavioralCloningSettings, ScheduleType from mlagents.trainers.torch.agent_action import AgentAction from mlagents.trainers.torch.action_log_probs import ActionLogProbs from mlagents.trainers.torch.utils import ModelUtils from mlagents.trainers.trajectory import ObsUtil from mlagents.trainers.buffer import AgentBuffer +from mlagents_envs import logging_util + +logger = logging_util.get_logger(__name__) class BCModule: @@ -33,13 +36,15 @@ def __init__( self._anneal_steps = settings.steps self.current_lr = policy_learning_rate * settings.strength - learning_rate_schedule: ScheduleType = ScheduleType.LINEAR if self._anneal_steps > 0 else ScheduleType.CONSTANT + learning_rate_schedule: ScheduleType = ( + ScheduleType.LINEAR if self._anneal_steps > 0 else ScheduleType.CONSTANT + ) self.decay_learning_rate = ModelUtils.DecayedValue( learning_rate_schedule, self.current_lr, 1e-10, self._anneal_steps ) params = self.policy.actor.parameters() self.optimizer = torch.optim.Adam(params, lr=self.current_lr) - _, self.demonstration_buffer = demo_to_buffer( + self._demo_manager = DemoManager( settings.demo_path, policy.sequence_length, policy.behavior_spec ) self.batch_size = ( @@ -47,7 +52,7 @@ def __init__( ) self.num_epoch = settings.num_epoch if settings.num_epoch else default_num_epoch self.n_sequences = max( - min(self.batch_size, self.demonstration_buffer.num_experiences) + min(self.batch_size, self._demo_manager.demo_buffer.num_experiences) // policy.sequence_length, 1, ) @@ -62,15 +67,23 @@ def update(self) -> Dict[str, np.ndarray]: :param max_batches: The maximum number of batches to use per update. :return: The loss of the update. """ - # Don't continue training if the learning rate has reached 0, to reduce training time. + num_new = self._demo_manager.refresh() + if num_new > 0: + # This is printed here so it is module-specific. + logger.info(f"Loaded {num_new} new demo samples for Behavior Cloning.") + # Check if we have no demos at all. + if len(self._demo_manager.demo_buffer.keys()) == 0: + return {} + + # Don't continue training if the learning rate has reached 0, to reduce training time. decay_lr = self.decay_learning_rate.get_value(self.policy.get_current_step()) if self.current_lr <= 1e-10: # Unlike in TF, this never actually reaches 0. return {"Losses/Pretraining Loss": 0} batch_losses = [] possible_demo_batches = ( - self.demonstration_buffer.num_experiences // self.n_sequences + self._demo_manager.demo_buffer.num_experiences // self.n_sequences ) possible_batches = possible_demo_batches @@ -78,7 +91,7 @@ def update(self) -> Dict[str, np.ndarray]: n_epoch = self.num_epoch for _ in range(n_epoch): - self.demonstration_buffer.shuffle( + self._demo_manager.demo_buffer.shuffle( sequence_length=self.policy.sequence_length ) if max_batches == 0: @@ -86,7 +99,7 @@ def update(self) -> Dict[str, np.ndarray]: else: num_batches = min(possible_batches, max_batches) for i in range(num_batches // self.policy.sequence_length): - demo_update_buffer = self.demonstration_buffer + demo_update_buffer = self._demo_manager.demo_buffer start = i * self.n_sequences * self.policy.sequence_length end = (i + 1) * self.n_sequences * self.policy.sequence_length mini_batch_demo = demo_update_buffer.make_mini_batch(start, end) diff --git a/ml-agents/mlagents/trainers/torch/components/reward_providers/gail_reward_provider.py b/ml-agents/mlagents/trainers/torch/components/reward_providers/gail_reward_provider.py index 031fd222110..fdc0c0cafe9 100644 --- a/ml-agents/mlagents/trainers/torch/components/reward_providers/gail_reward_provider.py +++ b/ml-agents/mlagents/trainers/torch/components/reward_providers/gail_reward_provider.py @@ -10,11 +10,11 @@ from mlagents_envs.base_env import BehaviorSpec from mlagents_envs import logging_util from mlagents.trainers.torch.utils import ModelUtils +from mlagents.trainers.demo_loader import DemoManager from mlagents.trainers.torch.agent_action import AgentAction from mlagents.trainers.torch.action_flattener import ActionFlattener from mlagents.trainers.torch.networks import NetworkBody from mlagents.trainers.torch.layers import linear_layer, Initialization -from mlagents.trainers.demo_loader import demo_to_buffer from mlagents.trainers.trajectory import ObsUtil logger = logging_util.get_logger(__name__) @@ -26,7 +26,7 @@ def __init__(self, specs: BehaviorSpec, settings: GAILSettings) -> None: self._ignore_done = False self._discriminator_network = DiscriminatorNetwork(specs, settings) self._discriminator_network.to(default_device()) - _, self._demo_buffer = demo_to_buffer( + self._demo_manager = DemoManager( settings.demo_path, 1, specs ) # This is supposed to be the sequence length but we do not have access here params = list(self._discriminator_network.parameters()) @@ -46,8 +46,16 @@ def evaluate(self, mini_batch: AgentBuffer) -> np.ndarray: ) def update(self, mini_batch: AgentBuffer) -> Dict[str, np.ndarray]: + num_new = self._demo_manager.refresh() + if num_new > 0: + # This is printed here so it is module-specific. + logger.info(f"Loaded {num_new} demo samples for GAIL.") - expert_batch = self._demo_buffer.sample_mini_batch( + # Check if we have no demos at all. + if len(self._demo_manager.demo_buffer.keys()) == 0: + return {} + + expert_batch = self._demo_manager.demo_buffer.sample_mini_batch( mini_batch.num_experiences, 1 ) self._discriminator_network.encoder.update_normalization(expert_batch)