Misc/veridream

hyperparams/human.yml

@@ -0,0 +1,7 @@
+# Space Engineers envs
+SE-WalkingTest-v1:
+  env_wrapper:
+    - utils.wrappers.HistoryWrapper:
+        horizon: 2
+  n_timesteps: !!float 2e6
+  policy: 'MlpPolicy'

hyperparams/tqc.yml

@@ -254,3 +254,97 @@ parking-v0:
     n_sampled_goal=4,
     max_episode_length=100
   )"
+
+# Space Engineers envs
+SE-Forward-v1: &defaults
+  env_wrapper:
+    - utils.wrappers.HistoryWrapper:
+        horizon: 2
+  vec_env_wrapper:
+    - utils.wrappers.VecForceResetWrapper
+  callback:
+    - utils.callbacks.ParallelTrainCallback:
+        gradient_steps: 400
+    # - utils.callbacks.StopTrainingOnMeanRewardThreshold:
+    #     reward_threshold: 250
+    #     verbose: 1
+  n_timesteps: !!float 5e6
+  policy: 'MlpPolicy'
+  learning_rate: !!float 7.3e-4
+  buffer_size: 100000
+  batch_size: 256
+  ent_coef: 'auto'
+  gamma: 0.98
+  tau: 0.05
+  # train_freq: [1, "episode"]
+  train_freq: 100
+  n_envs: 4
+  gradient_steps: -1
+  learning_starts: 800
+  use_sde: False
+  top_quantiles_to_drop_per_net: 2
+  policy_kwargs: "dict(net_arch=[256, 256], n_critics=2)"
+
+SE-Symmetric-v1:
+  <<: *defaults
+
+SE-Corrections-v1:
+  <<: *defaults
+
+SE-Generic-v1:
+  <<: *defaults
+  callback:
+    - utils.callbacks.ParallelTrainCallback:
+        gradient_steps: 400
+    # - utils.callbacks.StopTrainingOnMeanRewardThreshold:
+    #     reward_threshold: 250
+    #     verbose: 1
+
+SE-TurnLeft-v1:
+  <<: *defaults
+  callback:
+    - utils.callbacks.ParallelTrainCallback:
+        gradient_steps: 400
+    - utils.callbacks.StopTrainingOnMeanRewardThreshold:
+        reward_threshold: 250
+        verbose: 1
+
+SE-MultiTask-v1:
+  <<: *defaults
+  # policy: 'MixtureMlpPolicy'
+  learning_rate: !!float 7.3e-4
+  # gamma: 0.99
+  # tau: 0.005
+  buffer_size: 200000
+  callback:
+    - utils.callbacks.ParallelTrainCallback:
+        gradient_steps: 400
+  # policy_kwargs: "dict(net_arch=[400, 300], n_critics=2, n_additional_experts=2)"
+  policy_kwargs: "dict(net_arch=[256, 256], n_critics=5)"
+
+
+# ======== Real Robot envs ============
+
+WalkingBertSim-v1:
+  env_wrapper:
+    - utils.wrappers.HistoryWrapper:
+        horizon: 2
+  callback:
+    - utils.callbacks.ParallelTrainCallback:
+        gradient_steps: 400
+  n_timesteps: !!float 2e6
+  policy: 'MlpPolicy'
+  learning_rate: !!float 7.3e-4
+  buffer_size: 300000
+  batch_size: 256
+  ent_coef: 'auto'
+  gamma: 0.98
+  tau: 0.02
+  train_freq: [1, "episode"]
+  gradient_steps: -1
+  learning_starts: 1200
+  use_sde_at_warmup: True
+  use_sde: True
+  sde_sample_freq: 4
+  top_quantiles_to_drop_per_net: 2
+  policy_kwargs: "dict(log_std_init=-3, net_arch=[256, 256], n_critics=2)"

requirements.txt

@@ -1,8 +1,8 @@
 gym>=0.17,<0.20
-stable-baselines3[extra,tests,docs]>=1.3.1a8
-sb3-contrib>=1.3.1a7
-box2d-py==2.3.8
-pybullet
+# stable-baselines3[extra,tests,docs]>=1.3.0
+sb3-contrib>=1.3.0
+# box2d-py==2.3.8
+# pybullet
 gym-minigrid
 scikit-optimize
 optuna
@@ -11,7 +11,9 @@ seaborn
 pyyaml>=5.1
 cloudpickle>=1.5.0
 # tmp fix: ROM missing in newest release
-atari-py==0.2.6
+# atari-py==0.2.6
 plotly
-panda-gym==1.1.1 # tmp fix: until compatibility with panda-gym v2
-rliable>=1.0.5
+pygame
+# panda-gym>=1.1.1
+# rliable requires python 3.7+
+# rliable>=1.0.5

setup.cfg

@@ -20,6 +20,7 @@ per-file-ignores =
   ./scripts/all_plots.py:E501
   ./scripts/plot_train.py:E501
   ./scripts/plot_training_success.py:E501
+  ./utils/teleop.py:F405
 
 exclude =
     # No need to traverse our git directory

utils/callbacks.py

@@ -1,4 +1,5 @@
 import os
+import pickle
 import tempfile
 import time
 from copy import deepcopy
@@ -10,7 +11,8 @@
 from sb3_contrib import TQC
 from stable_baselines3 import SAC
 from stable_baselines3.common.callbacks import BaseCallback, EvalCallback
-from stable_baselines3.common.vec_env import VecEnv
+from stable_baselines3.common.utils import safe_mean
+from stable_baselines3.common.vec_env import VecEnv, sync_envs_normalization
 
 
 class TrialEvalCallback(EvalCallback):
@@ -129,6 +131,12 @@ def _init_callback(self) -> None:
 
         self.model.save(temp_file)
 
+        if self.model.get_vec_normalize_env() is not None:
+            temp_file_norm = os.path.join("logs", "vec_normalize.pkl")
+
+            with open(temp_file_norm, "wb") as file_handler:
+                pickle.dump(self.model.get_vec_normalize_env(), file_handler)
+
         # TODO: add support for other algorithms
         for model_class in [SAC, TQC]:
             if isinstance(self.model, model_class):
@@ -138,6 +146,11 @@ def _init_callback(self) -> None:
         assert self.model_class is not None, f"{self.model} is not supported for parallel training"
         self._model = self.model_class.load(temp_file)
 
+        if self.model.get_vec_normalize_env() is not None:
+            with open(temp_file_norm, "rb") as file_handler:
+                self._model._vec_normalize_env = pickle.load(file_handler)
+                self._model._vec_normalize_env.training = False
+
         self.batch_size = self._model.batch_size
 
         # Disable train method
@@ -182,6 +195,10 @@ def _on_rollout_end(self) -> None:
             self._model.replay_buffer = deepcopy(self.model.replay_buffer)
             self.model.set_parameters(deepcopy(self._model.get_parameters()))
             self.model.actor = self.model.policy.actor
+            # Sync VecNormalize
+            if self.model.get_vec_normalize_env() is not None:
+                sync_envs_normalization(self.model.get_vec_normalize_env(), self._model._vec_normalize_env)
+
             if self.num_timesteps >= self._model.learning_starts:
                 self.train()
             # Do not wait for the training loop to finish
@@ -193,3 +210,31 @@ def _on_training_end(self) -> None:
             if self.verbose > 0:
                 print("Waiting for training thread to terminate")
             self.process.join()
+
+
+class StopTrainingOnMeanRewardThreshold(BaseCallback):
+    """
+    Stop the training once a threshold in mean episodic reward
+    has been reached (i.e. when the model is good enough).
+
+    :param reward_threshold:  Minimum expected reward per episode
+        to stop training.
+    :param verbose:
+    """
+
+    def __init__(self, reward_threshold: float, verbose: int = 0):
+        super().__init__(verbose=verbose)
+        self.reward_threshold = reward_threshold
+
+    def _on_step(self) -> bool:
+        continue_training = True
+        if len(self.model.ep_info_buffer) > 0 and len(self.model.ep_info_buffer[0]) > 0:
+            mean_reward = safe_mean([ep_info["r"] for ep_info in self.model.ep_info_buffer])
+            # Convert np.bool_ to bool, otherwise callback() is False won't work
+            continue_training = bool(mean_reward < self.reward_threshold)
+            if self.verbose > 0 and not continue_training:
+                print(
+                    f"Stopping training because the mean reward {mean_reward:.2f} "
+                    f" is above the threshold {self.reward_threshold}"
+                )
+        return continue_training

utils/exp_manager.py

@@ -11,6 +11,7 @@
 import numpy as np
 import optuna
 import yaml
+import zmq
 from optuna.integration.skopt import SkoptSampler
 from optuna.pruners import BasePruner, MedianPruner, SuccessiveHalvingPruner
 from optuna.samplers import BaseSampler, RandomSampler, TPESampler
@@ -29,6 +30,7 @@
     DummyVecEnv,
     SubprocVecEnv,
     VecEnv,
+    VecEnvWrapper,
     VecFrameStack,
     VecNormalize,
     VecTransposeImage,
@@ -99,6 +101,7 @@ def __init__(
         self.env_wrapper = None
         self.frame_stack = None
         self.seed = seed
+        self.vec_env_wrapper = None
         self.optimization_log_path = optimization_log_path
 
         self.vec_env_class = {"dummy": DummyVecEnv, "subproc": SubprocVecEnv}[vec_env_type]
@@ -160,7 +163,7 @@ def setup_experiment(self) -> Optional[BaseAlgorithm]:
         :return: the initialized RL model
         """
         hyperparams, saved_hyperparams = self.read_hyperparameters()
-        hyperparams, self.env_wrapper, self.callbacks = self._preprocess_hyperparams(hyperparams)
+        hyperparams, self.env_wrapper, self.callbacks, self.vec_env_wrapper = self._preprocess_hyperparams(hyperparams)
 
         self.create_log_folder()
         self.create_callbacks()
@@ -200,12 +203,18 @@ def learn(self, model: BaseAlgorithm) -> None:
 
         try:
             model.learn(self.n_timesteps, **kwargs)
-        except KeyboardInterrupt:
+        except (KeyboardInterrupt, zmq.error.ZMQError):
             # this allows to save the model when interrupting training
             pass
         finally:
             # Release resources
             try:
+                # Hack for zmq on Windows to allow early termination
+                env_tmp = model.env
+                while isinstance(env_tmp, VecEnvWrapper):
+                    env_tmp = env_tmp.venv
+                env_tmp.waiting = False
+
                 model.env.close()
             except EOFError:
                 pass
@@ -310,7 +319,7 @@ def _preprocess_normalization(self, hyperparams: Dict[str, Any]) -> Dict[str, An
 
     def _preprocess_hyperparams(
         self, hyperparams: Dict[str, Any]
-    ) -> Tuple[Dict[str, Any], Optional[Callable], List[BaseCallback]]:
+    ) -> Tuple[Dict[str, Any], Optional[Callable], List[BaseCallback], Optional[Callable]]:
         self.n_envs = hyperparams.get("n_envs", 1)
 
         if self.verbose > 0:
@@ -354,12 +363,16 @@ def _preprocess_hyperparams(
         if "env_wrapper" in hyperparams.keys():
             del hyperparams["env_wrapper"]
 
+        vec_env_wrapper = get_wrapper_class(hyperparams, "vec_env_wrapper")
+        if "vec_env_wrapper" in hyperparams.keys():
+            del hyperparams["vec_env_wrapper"]
+
         callbacks = get_callback_list(hyperparams)
         if "callback" in hyperparams.keys():
             self.specified_callbacks = hyperparams["callback"]
             del hyperparams["callback"]
 
-        return hyperparams, env_wrapper, callbacks
+        return hyperparams, env_wrapper, callbacks, vec_env_wrapper
 
     def _preprocess_action_noise(
         self, hyperparams: Dict[str, Any], saved_hyperparams: Dict[str, Any], env: VecEnv
@@ -517,6 +530,9 @@ def create_envs(self, n_envs: int, eval_env: bool = False, no_log: bool = False)
             monitor_kwargs=monitor_kwargs,
         )
 
+        if self.vec_env_wrapper is not None:
+            env = self.vec_env_wrapper(env)
+
         # Wrap the env into a VecNormalize wrapper if needed
         # and load saved statistics when present
         env = self._maybe_normalize(env, eval_env)
@@ -653,9 +669,30 @@ def objective(self, trial: optuna.Trial) -> float:
         try:
             model.learn(self.n_timesteps, callback=callbacks)
             # Free memory
+            env_tmp = model.env
+            while isinstance(env_tmp, VecEnvWrapper):
+                env_tmp = env_tmp.venv
+            env_tmp.waiting = False
+
+            env_tmp = eval_env
+            while isinstance(env_tmp, VecEnvWrapper):
+                env_tmp = env_tmp.venv
+            env_tmp.waiting = False
+
             model.env.close()
             eval_env.close()
         except (AssertionError, ValueError) as e:
+            # Hack for zmq on Windows to allow early termination
+            env_tmp = model.env
+            while isinstance(env_tmp, VecEnvWrapper):
+                env_tmp = env_tmp.venv
+            env_tmp.waiting = False
+
+            env_tmp = eval_env
+            while isinstance(env_tmp, VecEnvWrapper):
+                env_tmp = env_tmp.venv
+            env_tmp.waiting = False
+
             # Sometimes, random hyperparams can generate NaN
             # Free memory
             model.env.close()