Repository: DLR-RM/rl-baselines3-zoo (language: python)

Issue:
[feature request] make rl-baselines3-zoo a python package
If you have any questions, feel free to create an issue with the tag [question].  
If you wish to suggest an enhancement or feature request, add the tag [feature request].  
If you are submitting a bug report, please fill in the following details.  

**Describe the bug**
I have a project that constructs a ad-hoc module around a trained RL policy. With that being said, I also have my own code space and I would just simply want to load trained policy available in `rl-trained-agents`. The `enjoy.py` in `rl-baselines3-zoo` defined some modules under `utils` folder to assist the loading and playing procedures. However, since the `rl-baselines3-zoo` is not a python package, I cannot call those utility modules directly outside the `rl-baselines3-zoo` folder. I think making the `rl-baselines3-zoo` as an package that can be installed and called from outside code would significantly increase the code reusability.

**Code example**
N/A

**System Info**
Describe the characteristic of your environment:
 * pip install -e .[extra]
 * ThinkPad T480s with MX150
 * Python 3.8.10
 * PyTorch 1.11
 * Relevant libraries installed with pip and everything works fine.

**Additional context**
Add any other context about the problem here.

RL Zoo CLI
It would be nice to have a cli that is available everywhere (with completion too), something like:
```bash
sb3 train --algo ppo --env CartPole-v1  
sb3 train --algo ppo --env CartPole-v1  -i custom_params.yaml # allow overwriting hyperparams file with a custom one
sb3 enjoy --algo sac --env Pendulum-v0
sb3 plot_train -a a2c -e Breakout
sb3 plot_eval ...
sb3 plot_from_file
```


Discussion / hints:
Hello,

so that question was already asked in https://github.com/araffin/rl-baselines-zoo/issues/107 and https://github.com/DLR-RM/rl-baselines3-zoo/issues/53

I guess what you would like to access is the loading utils from `utils/` folder? (that is also in charge of wrapping the env).

I think this is good enough use-case to at least add a small `setup.py` in order to load pretrained agents easily.

> However, since the rl-baselines3-zoo is not a python package, I cannot call those utility modules directly outside the rl-baselines3-zoo folder.

well, it is possible by adding those file to the path, but not pretty yes.

Side note: we have now huggingface hub support, which should facilitate export and loading of pre-trained agent, see https://huggingface.co/sb3/ppo-MountainCarContinuous-v0 (with the zoo) or https://huggingface.co/araffin/a2c-LunarLander-v2 (independent of the zoo)
Hey @araffin! I've recently done simple `console_scripts` entry point for the package (to help running scripts from folders with custom code). I can polish it and submit PR.  It seems like entry point is the most reliable way for this to work on different platforms. Though a few caveats:

* "everywhere" is typically not desirable because of dependencies management. `console_scripts` could work globally if the package was installed globally. But I assume in most cases packages are installed into virtual environment, and in this case `sb3` would be available only within active env.
* Completion is somewhat complicated functionality to get reliable on multiple platforms/terminals. Usually requires update for terminal init files (like `~/.bash_profile` and so on). But I can check what is a common approach for doing it nowadays.
Hello,
thanks for the proposal =)

>  and in this case sb3 would be available only within active env.

yes, that's what I meant by "everywhere"  (more in the sense "from every folder rather than from every env")

> Completion is somewhat complicated functionality to get reliable on multiple platforms/terminals

that's true, but I had in mind to provide only minimal completion (different first level command), at least for bash/zsh users.
I did that in the past for a wrapper around conan that tries to mimic conda env.

> I've recently done simple console_scripts entry point for the package 

You can submit a draft PR yes, but I'm not sure what you mean exactly by "console scripts".
(I also have private code of a cli I created for another purpose (but not for the zoo))

`console_scripts` is one the functionality `setuptools` provides for automatically generating script wrappers for modules within the packages, link to their [docs](https://setuptools.pypa.io/en/latest/userguide/entry_point.html#entry-points).

Fix patch:
diff --git a/.coveragerc b/.coveragerc
index 820175c82..785267de2 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -2,7 +2,7 @@
 branch = False
 omit =
     tests/*
-    utils/plot.py
+    rl_zoo/utils/plot.py
 
 [report]
 exclude_lines =
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0e78317f9..d773bcee9 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -36,6 +36,7 @@ jobs:
         pip install opencv-python-headless
         # install parking-env to test HER (pinned so it works with gym 0.21)
         pip install highway-env==1.5.0
+        pip install -e .
     - name: Type check
       run: |
         make type
diff --git a/.github/workflows/trained_agents.yml b/.github/workflows/trained_agents.yml
index a047e9861..7647a5967 100644
--- a/.github/workflows/trained_agents.yml
+++ b/.github/workflows/trained_agents.yml
@@ -38,6 +38,7 @@ jobs:
         pip install highway-env==1.5.0
         # Add support for pickle5 protocol
         pip install pickle5
+        pip install -e .
     - name: Check trained agents
       run: |
         make check-trained-agents
diff --git a/.gitignore b/.gitignore
index bff14fe14..87c26d1a2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,3 +18,12 @@ runs
 hub
 *.mp4
 *.json
+
+# Setuptools distribution and build folders.
+/dist/
+/build
+keys/
+*.egg-info
+.cache
+*.lprof
+*.prof
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 531d6665e..ef76b847c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,18 @@
+## Release 1.6.2 (2022-10-02)
+
+### Breaking Changes
+- RL Zoo is now a python package
+- low pass filter was removed
+
+### New Features
+- RL Zoo cli: `rl_zoo train` and `rl_zoo enjoy`
+
+### Bug fixes
+
+### Documentation
+
+### Other
+
 ## Release 1.6.1 (2022-09-30)
 
 **Progress bar and custom yaml file**
diff --git a/Makefile b/Makefile
index bb8788e6c..8f20105df 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-LINT_PATHS = *.py tests/ scripts/ utils/
+LINT_PATHS = *.py tests/ scripts/ rl_zoo/
 
 # Run pytest and coverage report
 pytest:
@@ -10,7 +10,7 @@ check-trained-agents:
 
 # Type check
 type:
-	pytype -j auto ${LINT_PATHS}
+	pytype -j auto rl_zoo/ tests/ scripts/ -d import-error
 
 lint:
 	# stop the build if there are Python syntax errors or undefined names
@@ -42,4 +42,16 @@ docker-cpu:
 docker-gpu:
 	USE_GPU=True ./scripts/build_docker.sh
 
-.PHONY: docker lint type pytest
+# PyPi package release
+release:
+	python setup.py sdist
+	python setup.py bdist_wheel
+	twine upload dist/*
+
+# Test PyPi package release
+test-release:
+	python setup.py sdist
+	python setup.py bdist_wheel
+	twine upload --repository-url https://test.pypi.org/legacy/ dist/*
+
+.PHONY: lint format check-codestyle commit-checks doc spelling  docker type pytest
diff --git a/enjoy.py b/enjoy.py
index d94fa03fa..926b70141 100644
--- a/enjoy.py
+++ b/enjoy.py
@@ -1,274 +1,4 @@
-import argparse
-import importlib
-import os
-import sys
-
-import numpy as np
-import torch as th
-import yaml
-from huggingface_sb3 import EnvironmentName
-from stable_baselines3.common.utils import set_random_seed
-
-import utils.import_envs  # noqa: F401 pylint: disable=unused-import
-from utils import ALGOS, create_test_env, get_saved_hyperparams
-from utils.callbacks import tqdm
-from utils.exp_manager import ExperimentManager
-from utils.load_from_hub import download_from_hub
-from utils.utils import StoreDict, get_model_path
-
-
-def main():  # noqa: C901
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--env", help="environment ID", type=EnvironmentName, default="CartPole-v1")
-    parser.add_argument("-f", "--folder", help="Log folder", type=str, default="rl-trained-agents")
-    parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False, choices=list(ALGOS.keys()))
-    parser.add_argument("-n", "--n-timesteps", help="number of timesteps", default=1000, type=int)
-    parser.add_argument("--num-threads", help="Number of threads for PyTorch (-1 to use default)", default=-1, type=int)
-    parser.add_argument("--n-envs", help="number of environments", default=1, type=int)
-    parser.add_argument("--exp-id", help="Experiment ID (default: 0: latest, -1: no exp folder)", default=0, type=int)
-    parser.add_argument("--verbose", help="Verbose mode (0: no output, 1: INFO)", default=1, type=int)
-    parser.add_argument(
-        "--no-render", action="store_true", default=False, help="Do not render the environment (useful for tests)"
-    )
-    parser.add_argument("--deterministic", action="store_true", default=False, help="Use deterministic actions")
-    parser.add_argument("--device", help="PyTorch device to be use (ex: cpu, cuda...)", default="auto", type=str)
-    parser.add_argument(
-        "--load-best", action="store_true", default=False, help="Load best model instead of last model if available"
-    )
-    parser.add_argument(
-        "--load-checkpoint",
-        type=int,
-        help="Load checkpoint instead of last model if available, "
-        "you must pass the number of timesteps corresponding to it",
-    )
-    parser.add_argument(
-        "--load-last-checkpoint",
-        action="store_true",
-        default=False,
-        help="Load last checkpoint instead of last model if available",
-    )
-    parser.add_argument("--stochastic", action="store_true", default=False, help="Use stochastic actions")
-    parser.add_argument(
-        "--norm-reward", action="store_true", default=False, help="Normalize reward if applicable (trained with VecNormalize)"
-    )
-    parser.add_argument("--seed", help="Random generator seed", type=int, default=0)
-    parser.add_argument("--reward-log", help="Where to log reward", default="", type=str)
-    parser.add_argument(
-        "--gym-packages",
-        type=str,
-        nargs="+",
-        default=[],
-        help="Additional external Gym environment package modules to import (e.g. gym_minigrid)",
-    )
-    parser.add_argument(
-        "--env-kwargs", type=str, nargs="+", action=StoreDict, help="Optional keyword argument to pass to the env constructor"
-    )
-    parser.add_argument(
-        "--custom-objects", action="store_true", default=False, help="Use custom objects to solve loading issues"
-    )
-    parser.add_argument(
-        "-P",
-        "--progress",
-        action="store_true",
-        default=False,
-        help="if toggled, display a progress bar using tqdm and rich",
-    )
-    args = parser.parse_args()
-
-    # Going through custom gym packages to let them register in the global registory
-    for env_module in args.gym_packages:
-        importlib.import_module(env_module)
-
-    env_name: EnvironmentName = args.env
-    algo = args.algo
-    folder = args.folder
-
-    try:
-        _, model_path, log_path = get_model_path(
-            args.exp_id,
-            folder,
-            algo,
-            env_name,
-            args.load_best,
-            args.load_checkpoint,
-            args.load_last_checkpoint,
-        )
-    except (AssertionError, ValueError) as e:
-        # Special case for rl-trained agents
-        # auto-download from the hub
-        if "rl-trained-agents" not in folder:
-            raise e
-        else:
-            print("Pretrained model not found, trying to download it from sb3 Huggingface hub: https://huggingface.co/sb3")
-            # Auto-download
-            download_from_hub(
-                algo=algo,
-                env_name=env_name,
-                exp_id=args.exp_id,
-                folder=folder,
-                organization="sb3",
-                repo_name=None,
-                force=False,
-            )
-            # Try again
-            _, model_path, log_path = get_model_path(
-                args.exp_id,
-                folder,
-                algo,
-                env_name,
-                args.load_best,
-                args.load_checkpoint,
-                args.load_last_checkpoint,
-            )
-
-    print(f"Loading {model_path}")
-
-    # Off-policy algorithm only support one env for now
-    off_policy_algos = ["qrdqn", "dqn", "ddpg", "sac", "her", "td3", "tqc"]
-
-    if algo in off_policy_algos:
-        args.n_envs = 1
-
-    set_random_seed(args.seed)
-
-    if args.num_threads > 0:
-        if args.verbose > 1:
-            print(f"Setting torch.num_threads to {args.num_threads}")
-        th.set_num_threads(args.num_threads)
-
-    is_atari = ExperimentManager.is_atari(env_name.gym_id)
-
-    stats_path = os.path.join(log_path, env_name)
-    hyperparams, stats_path = get_saved_hyperparams(stats_path, norm_reward=args.norm_reward, test_mode=True)
-
-    # load env_kwargs if existing
-    env_kwargs = {}
-    args_path = os.path.join(log_path, env_name, "args.yml")
-    if os.path.isfile(args_path):
-        with open(args_path) as f:
-            loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)  # pytype: disable=module-attr
-            if loaded_args["env_kwargs"] is not None:
-                env_kwargs = loaded_args["env_kwargs"]
-    # overwrite with command line arguments
-    if args.env_kwargs is not None:
-        env_kwargs.update(args.env_kwargs)
-
-    log_dir = args.reward_log if args.reward_log != "" else None
-
-    env = create_test_env(
-        env_name.gym_id,
-        n_envs=args.n_envs,
-        stats_path=stats_path,
-        seed=args.seed,
-        log_dir=log_dir,
-        should_render=not args.no_render,
-        hyperparams=hyperparams,
-        env_kwargs=env_kwargs,
-    )
-
-    kwargs = dict(seed=args.seed)
-    if algo in off_policy_algos:
-        # Dummy buffer size as we don't need memory to enjoy the trained agent
-        kwargs.update(dict(buffer_size=1))
-        # Hack due to breaking change in v1.6
-        # handle_timeout_termination cannot be at the same time
-        # with optimize_memory_usage
-        if "optimize_memory_usage" in hyperparams:
-            kwargs.update(optimize_memory_usage=False)
-
-    # Check if we are running python 3.8+
-    # we need to patch saved model under python 3.6/3.7 to load them
-    newer_python_version = sys.version_info.major == 3 and sys.version_info.minor >= 8
-
-    custom_objects = {}
-    if newer_python_version or args.custom_objects:
-        custom_objects = {
-            "learning_rate": 0.0,
-            "lr_schedule": lambda _: 0.0,
-            "clip_range": lambda _: 0.0,
-        }
-
-    model = ALGOS[algo].load(model_path, env=env, custom_objects=custom_objects, device=args.device, **kwargs)
-
-    obs = env.reset()
-
-    # Deterministic by default except for atari games
-    stochastic = args.stochastic or is_atari and not args.deterministic
-    deterministic = not stochastic
-
-    episode_reward = 0.0
-    episode_rewards, episode_lengths = [], []
-    ep_len = 0
-    # For HER, monitor success rate
-    successes = []
-    lstm_states = None
-    episode_start = np.ones((env.num_envs,), dtype=bool)
-
-    generator = range(args.n_timesteps)
-    if args.progress:
-        generator = tqdm(generator)
-
-    try:
-        for _ in generator:
-            action, lstm_states = model.predict(
-                obs,
-                state=lstm_states,
-                episode_start=episode_start,
-                deterministic=deterministic,
-            )
-            obs, reward, done, infos = env.step(action)
-
-            episode_start = done
-
-            if not args.no_render:
-                env.render("human")
-
-            episode_reward += reward[0]
-            ep_len += 1
-
-            if args.n_envs == 1:
-                # For atari the return reward is not the atari score
-                # so we have to get it from the infos dict
-                if is_atari and infos is not None and args.verbose >= 1:
-                    episode_infos = infos[0].get("episode")
-                    if episode_infos is not None:
-                        print(f"Atari Episode Score: {episode_infos['r']:.2f}")
-                        print("Atari Episode Length", episode_infos["l"])
-
-                if done and not is_atari and args.verbose > 0:
-                    # NOTE: for env using VecNormalize, the mean reward
-                    # is a normalized reward when `--norm_reward` flag is passed
-                    print(f"Episode Reward: {episode_reward:.2f}")
-                    print("Episode Length", ep_len)
-                    episode_rewards.append(episode_reward)
-                    episode_lengths.append(ep_len)
-                    episode_reward = 0.0
-                    ep_len = 0
-
-                # Reset also when the goal is achieved when using HER
-                if done and infos[0].get("is_success") is not None:
-                    if args.verbose > 1:
-                        print("Success?", infos[0].get("is_success", False))
-
-                    if infos[0].get("is_success") is not None:
-                        successes.append(infos[0].get("is_success", False))
-                        episode_reward, ep_len = 0.0, 0
-
-    except KeyboardInterrupt:
-        pass
-
-    if args.verbose > 0 and len(successes) > 0:
-        print(f"Success rate: {100 * np.mean(successes):.2f}%")
-
-    if args.verbose > 0 and len(episode_rewards) > 0:
-        print(f"{len(episode_rewards)} Episodes")
-        print(f"Mean reward: {np.mean(episode_rewards):.2f} +/- {np.std(episode_rewards):.2f}")
-
-    if args.verbose > 0 and len(episode_lengths) > 0:
-        print(f"Mean episode length: {np.mean(episode_lengths):.2f} +/- {np.std(episode_lengths):.2f}")
-
-    env.close()
-
+from rl_zoo.enjoy import enjoy
 
 if __name__ == "__main__":
-    main()
+    enjoy()
diff --git a/hyperparams/her.yml b/hyperparams/her.yml
index 9de53054b..45c0710fe 100644
--- a/hyperparams/her.yml
+++ b/hyperparams/her.yml
@@ -59,7 +59,7 @@ FetchSlide-v1:
 FetchPickAndPlace-v1:
   env_wrapper:
     - sb3_contrib.common.wrappers.TimeFeatureWrapper
-    # - utils.wrappers.DoneOnSuccessWrapper:
+    # - rl_zoo.wrappers.DoneOnSuccessWrapper:
     #     reward_offset: 0
     #     n_successes: 4
     # - stable_baselines3.common.monitor.Monitor
@@ -96,7 +96,7 @@ FetchReach-v1:
 NeckGoalEnvRelativeSparse-v2:
   model_class: 'sac'
   # env_wrapper:
-  #   - utils.wrappers.HistoryWrapper:
+  #   - rl_zoo.wrappers.HistoryWrapper:
   #       horizon: 2
   #   - sb3_contrib.common.wrappers.TimeFeatureWrapper
   n_timesteps: !!float 1e6
@@ -122,7 +122,7 @@ NeckGoalEnvRelativeSparse-v2:
 NeckGoalEnvRelativeDense-v2:
   model_class: 'sac'
   env_wrapper:
-    - utils.wrappers.HistoryWrapperObsDict:
+    - rl_zoo.wrappers.HistoryWrapperObsDict:
         horizon: 2
   #   - sb3_contrib.common.wrappers.TimeFeatureWrapper
   n_timesteps: !!float 1e6
diff --git a/hyperparams/ppo.yml b/hyperparams/ppo.yml
index 1532953d5..26d90d5f1 100644
--- a/hyperparams/ppo.yml
+++ b/hyperparams/ppo.yml
@@ -319,7 +319,7 @@ MiniGrid-FourRooms-v0:
 
 CarRacing-v0:
   env_wrapper:
-    - utils.wrappers.FrameSkip:
+    - rl_zoo.wrappers.FrameSkip:
         skip: 2
     - gym.wrappers.resize_observation.ResizeObservation:
         shape: 64
diff --git a/hyperparams/ppo_lstm.yml b/hyperparams/ppo_lstm.yml
index 52fffe32f..4043d7d4c 100644
--- a/hyperparams/ppo_lstm.yml
+++ b/hyperparams/ppo_lstm.yml
@@ -132,7 +132,7 @@ BipedalWalker-v3:
 # TO BE TUNED
 BipedalWalkerHardcore-v3:
   # env_wrapper:
-  #   - utils.wrappers.FrameSkip:
+  #   - rl_zoo.wrappers.FrameSkip:
   #       skip: 2
   normalize: true
   n_envs: 32
@@ -285,7 +285,7 @@ InvertedPendulumSwingupBulletEnv-v0:
 
 CarRacing-v0:
   env_wrapper:
-    # - utils.wrappers.FrameSkip:
+    # - rl_zoo.wrappers.FrameSkip:
     #     skip: 2
     - gym.wrappers.resize_observation.ResizeObservation:
         shape: 64
diff --git a/hyperparams/sac.yml b/hyperparams/sac.yml
index 1d3b3e02c..d4475fdce 100644
--- a/hyperparams/sac.yml
+++ b/hyperparams/sac.yml
@@ -16,7 +16,7 @@ MountainCarContinuous-v0:
 
 Pendulum-v1:
   # callback:
-  #   - utils.callbacks.ParallelTrainCallback
+  #   - rl_zoo.callbacks.ParallelTrainCallback
   n_timesteps: 20000
   policy: 'MlpPolicy'
   learning_rate: !!float 1e-3
@@ -74,9 +74,9 @@ BipedalWalkerHardcore-v3:
 HalfCheetahBulletEnv-v0: &pybullet-defaults
   # env_wrapper:
   #   - sb3_contrib.common.wrappers.TimeFeatureWrapper
-  #   - utils.wrappers.DelayedRewardWrapper:
+  #   - rl_zoo.wrappers.DelayedRewardWrapper:
   #       delay: 10
-  #   - utils.wrappers.HistoryWrapper:
+  #   - rl_zoo.wrappers.HistoryWrapper:
   #       horizon: 10
   n_timesteps: !!float 1e6
   policy: 'MlpPolicy'
@@ -163,12 +163,12 @@ MinitaurBulletDuckEnv-v0:
 # To be tuned
 CarRacing-v0:
   env_wrapper:
-    - utils.wrappers.FrameSkip:
+    - rl_zoo.wrappers.FrameSkip:
         skip: 2
     # wrapper from https://github.com/araffin/aae-train-donkeycar
     - ae.wrapper.AutoencoderWrapper:
         ae_path: "logs/car_racing_rgb_160.pkl"
-    - utils.wrappers.HistoryWrapper:
+    - rl_zoo.wrappers.HistoryWrapper:
         horizon: 2
   # frame_stack: 4
   normalize: True
@@ -238,7 +238,7 @@ donkey-generated-track-v0:
   env_wrapper:
     - gym.wrappers.time_limit.TimeLimit:
         max_episode_steps: 500
-    - utils.wrappers.HistoryWrapper:
+    - rl_zoo.wrappers.HistoryWrapper:
         horizon: 5
   n_timesteps: !!float 1e6
   policy: 'MlpPolicy'
@@ -262,9 +262,9 @@ donkey-generated-track-v0:
 NeckEnvRelative-v2:
   <<: *pybullet-defaults
   env_wrapper:
-    - utils.wrappers.HistoryWrapper:
+    - rl_zoo.wrappers.HistoryWrapper:
         horizon: 2
-    # - utils.wrappers.LowPassFilterWrapper:
+    # - rl_zoo.wrappers.LowPassFilterWrapper:
     #     freq: 2.0
     #     df: 25.0
   n_timesteps: !!float 1e6
diff --git a/hyperparams/tqc.yml b/hyperparams/tqc.yml
index e734cc48b..720eb91bf 100644
--- a/hyperparams/tqc.yml
+++ b/hyperparams/tqc.yml
@@ -258,12 +258,12 @@ parking-v0:
 # Tuned
 CarRacing-v0:
   env_wrapper:
-    - utils.wrappers.FrameSkip:
+    - rl_zoo.wrappers.FrameSkip:
         skip: 2
     # wrapper from https://github.com/araffin/aae-train-donkeycar
     - ae.wrapper.AutoencoderWrapper:
         ae_path: "logs/car_racing_rgb_160.pkl"
-    - utils.wrappers.HistoryWrapper:
+    - rl_zoo.wrappers.HistoryWrapper:
         horizon: 2
   # frame_stack: 4
   normalize: True
@@ -280,7 +280,7 @@ RocketLander-v0:
   n_timesteps: !!float 3e6
   policy: 'MlpPolicy'
   env_wrapper:
-    - utils.wrappers.FrameSkip:
+    - rl_zoo.wrappers.FrameSkip:
         skip: 4
-    - utils.wrappers.HistoryWrapper:
+    - rl_zoo.wrappers.HistoryWrapper:
         horizon: 2
diff --git a/requirements.txt b/requirements.txt
index 2d24dad9e..2739c4914 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -18,3 +18,4 @@ huggingface_sb3>=2.2.1, <3.*
 seaborn
 tqdm
 rich
+importlib-metadata~=4.13 # flake8 not compatible with importlib-metadata>5.0
diff --git a/rl_zoo/__init__.py b/rl_zoo/__init__.py
new file mode 100644
index 000000000..7df789c63
--- /dev/null
+++ b/rl_zoo/__init__.py
@@ -0,0 +1,16 @@
+import os
+
+from rl_zoo.utils import (
+    ALGOS,
+    create_test_env,
+    get_latest_run_id,
+    get_saved_hyperparams,
+    get_trained_models,
+    get_wrapper_class,
+    linear_schedule,
+)
+
+# Read version from file
+version_file = os.path.join(os.path.dirname(__file__), "version.txt")
+with open(version_file) as file_handler:
+    __version__ = file_handler.read().strip()
diff --git a/utils/benchmark.py b/rl_zoo/benchmark.py
similarity index 98%
rename from utils/benchmark.py
rename to rl_zoo/benchmark.py
index b53430fb2..01626fcec 100644
--- a/utils/benchmark.py
+++ b/rl_zoo/benchmark.py
@@ -9,7 +9,7 @@
 import pytablewriter
 from stable_baselines3.common.results_plotter import load_results, ts2xy
 
-from utils.utils import get_hf_trained_models, get_latest_run_id, get_saved_hyperparams, get_trained_models
+from rl_zoo.utils import get_hf_trained_models, get_latest_run_id, get_saved_hyperparams, get_trained_models
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--log-dir", help="Root log folder", default="rl-trained-agents/", type=str)
diff --git a/utils/callbacks.py b/rl_zoo/callbacks.py
similarity index 100%
rename from utils/callbacks.py
rename to rl_zoo/callbacks.py
diff --git a/rl_zoo/cli.py b/rl_zoo/cli.py
new file mode 100644
index 000000000..5902224f3
--- /dev/null
+++ b/rl_zoo/cli.py
@@ -0,0 +1,15 @@
+import sys
+
+from rl_zoo.enjoy import enjoy
+from rl_zoo.train import train
+
+
+def main():
+    script_name = sys.argv[1]
+    # Remove script name
+    del sys.argv[1]
+    # Execute known script
+    {
+        "train": train,
+        "enjoy": enjoy,
+    }[script_name]()
diff --git a/rl_zoo/enjoy.py b/rl_zoo/enjoy.py
new file mode 100644
index 000000000..9fe0b5218
--- /dev/null
+++ b/rl_zoo/enjoy.py
@@ -0,0 +1,274 @@
+import argparse
+import importlib
+import os
+import sys
+
+import numpy as np
+import torch as th
+import yaml
+from huggingface_sb3 import EnvironmentName
+from stable_baselines3.common.utils import set_random_seed
+
+import rl_zoo.import_envs  # noqa: F401 pylint: disable=unused-import
+from rl_zoo import ALGOS, create_test_env, get_saved_hyperparams
+from rl_zoo.callbacks import tqdm
+from rl_zoo.exp_manager import ExperimentManager
+from rl_zoo.load_from_hub import download_from_hub
+from rl_zoo.utils import StoreDict, get_model_path
+
+
+def enjoy():  # noqa: C901
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env", help="environment ID", type=EnvironmentName, default="CartPole-v1")
+    parser.add_argument("-f", "--folder", help="Log folder", type=str, default="rl-trained-agents")
+    parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False, choices=list(ALGOS.keys()))
+    parser.add_argument("-n", "--n-timesteps", help="number of timesteps", default=1000, type=int)
+    parser.add_argument("--num-threads", help="Number of threads for PyTorch (-1 to use default)", default=-1, type=int)
+    parser.add_argument("--n-envs", help="number of environments", default=1, type=int)
+    parser.add_argument("--exp-id", help="Experiment ID (default: 0: latest, -1: no exp folder)", default=0, type=int)
+    parser.add_argument("--verbose", help="Verbose mode (0: no output, 1: INFO)", default=1, type=int)
+    parser.add_argument(
+        "--no-render", action="store_true", default=False, help="Do not render the environment (useful for tests)"
+    )
+    parser.add_argument("--deterministic", action="store_true", default=False, help="Use deterministic actions")
+    parser.add_argument("--device", help="PyTorch device to be use (ex: cpu, cuda...)", default="auto", type=str)
+    parser.add_argument(
+        "--load-best", action="store_true", default=False, help="Load best model instead of last model if available"
+    )
+    parser.add_argument(
+        "--load-checkpoint",
+        type=int,
+        help="Load checkpoint instead of last model if available, "
+        "you must pass the number of timesteps corresponding to it",
+    )
+    parser.add_argument(
+        "--load-last-checkpoint",
+        action="store_true",
+        default=False,
+        help="Load last checkpoint instead of last model if available",
+    )
+    parser.add_argument("--stochastic", action="store_true", default=False, help="Use stochastic actions")
+    parser.add_argument(
+        "--norm-reward", action="store_true", default=False, help="Normalize reward if applicable (trained with VecNormalize)"
+    )
+    parser.add_argument("--seed", help="Random generator seed", type=int, default=0)
+    parser.add_argument("--reward-log", help="Where to log reward", default="", type=str)
+    parser.add_argument(
+        "--gym-packages",
+        type=str,
+        nargs="+",
+        default=[],
+        help="Additional external Gym environment package modules to import (e.g. gym_minigrid)",
+    )
+    parser.add_argument(
+        "--env-kwargs", type=str, nargs="+", action=StoreDict, help="Optional keyword argument to pass to the env constructor"
+    )
+    parser.add_argument(
+        "--custom-objects", action="store_true", default=False, help="Use custom objects to solve loading issues"
+    )
+    parser.add_argument(
+        "-P",
+        "--progress",
+        action="store_true",
+        default=False,
+        help="if toggled, display a progress bar using tqdm and rich",
+    )
+    args = parser.parse_args()
+
+    # Going through custom gym packages to let them register in the global registory
+    for env_module in args.gym_packages:
+        importlib.import_module(env_module)
+
+    env_name: EnvironmentName = args.env
+    algo = args.algo
+    folder = args.folder
+
+    try:
+        _, model_path, log_path = get_model_path(
+            args.exp_id,
+            folder,
+            algo,
+            env_name,
+            args.load_best,
+            args.load_checkpoint,
+            args.load_last_checkpoint,
+        )
+    except (AssertionError, ValueError) as e:
+        # Special case for rl-trained agents
+        # auto-download from the hub
+        if "rl-trained-agents" not in folder:
+            raise e
+        else:
+            print("Pretrained model not found, trying to download it from sb3 Huggingface hub: https://huggingface.co/sb3")
+            # Auto-download
+            download_from_hub(
+                algo=algo,
+                env_name=env_name,
+                exp_id=args.exp_id,
+                folder=folder,
+                organization="sb3",
+                repo_name=None,
+                force=False,
+            )
+            # Try again
+            _, model_path, log_path = get_model_path(
+                args.exp_id,
+                folder,
+                algo,
+                env_name,
+                args.load_best,
+                args.load_checkpoint,
+                args.load_last_checkpoint,
+            )
+
+    print(f"Loading {model_path}")
+
+    # Off-policy algorithm only support one env for now
+    off_policy_algos = ["qrdqn", "dqn", "ddpg", "sac", "her", "td3", "tqc"]
+
+    if algo in off_policy_algos:
+        args.n_envs = 1
+
+    set_random_seed(args.seed)
+
+    if args.num_threads > 0:
+        if args.verbose > 1:
+            print(f"Setting torch.num_threads to {args.num_threads}")
+        th.set_num_threads(args.num_threads)
+
+    is_atari = ExperimentManager.is_atari(env_name.gym_id)
+
+    stats_path = os.path.join(log_path, env_name)
+    hyperparams, stats_path = get_saved_hyperparams(stats_path, norm_reward=args.norm_reward, test_mode=True)
+
+    # load env_kwargs if existing
+    env_kwargs = {}
+    args_path = os.path.join(log_path, env_name, "args.yml")
+    if os.path.isfile(args_path):
+        with open(args_path) as f:
+            loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)  # pytype: disable=module-attr
+            if loaded_args["env_kwargs"] is not None:
+                env_kwargs = loaded_args["env_kwargs"]
+    # overwrite with command line arguments
+    if args.env_kwargs is not None:
+        env_kwargs.update(args.env_kwargs)
+
+    log_dir = args.reward_log if args.reward_log != "" else None
+
+    env = create_test_env(
+        env_name.gym_id,
+        n_envs=args.n_envs,
+        stats_path=stats_path,
+        seed=args.seed,
+        log_dir=log_dir,
+        should_render=not args.no_render,
+        hyperparams=hyperparams,
+        env_kwargs=env_kwargs,
+    )
+
+    kwargs = dict(seed=args.seed)
+    if algo in off_policy_algos:
+        # Dummy buffer size as we don't need memory to enjoy the trained agent
+        kwargs.update(dict(buffer_size=1))
+        # Hack due to breaking change in v1.6
+        # handle_timeout_termination cannot be at the same time
+        # with optimize_memory_usage
+        if "optimize_memory_usage" in hyperparams:
+            kwargs.update(optimize_memory_usage=False)
+
+    # Check if we are running python 3.8+
+    # we need to patch saved model under python 3.6/3.7 to load them
+    newer_python_version = sys.version_info.major == 3 and sys.version_info.minor >= 8
+
+    custom_objects = {}
+    if newer_python_version or args.custom_objects:
+        custom_objects = {
+            "learning_rate": 0.0,
+            "lr_schedule": lambda _: 0.0,
+            "clip_range": lambda _: 0.0,
+        }
+
+    model = ALGOS[algo].load(model_path, env=env, custom_objects=custom_objects, device=args.device, **kwargs)
+
+    obs = env.reset()
+
+    # Deterministic by default except for atari games
+    stochastic = args.stochastic or is_atari and not args.deterministic
+    deterministic = not stochastic
+
+    episode_reward = 0.0
+    episode_rewards, episode_lengths = [], []
+    ep_len = 0
+    # For HER, monitor success rate
+    successes = []
+    lstm_states = None
+    episode_start = np.ones((env.num_envs,), dtype=bool)
+
+    generator = range(args.n_timesteps)
+    if args.progress:
+        generator = tqdm(generator)
+
+    try:
+        for _ in generator:
+            action, lstm_states = model.predict(
+                obs,
+                state=lstm_states,
+                episode_start=episode_start,
+                deterministic=deterministic,
+            )
+            obs, reward, done, infos = env.step(action)
+
+            episode_start = done
+
+            if not args.no_render:
+                env.render("human")
+
+            episode_reward += reward[0]
+            ep_len += 1
+
+            if args.n_envs == 1:
+                # For atari the return reward is not the atari score
+                # so we have to get it from the infos dict
+                if is_atari and infos is not None and args.verbose >= 1:
+                    episode_infos = infos[0].get("episode")
+                    if episode_infos is not None:
+                        print(f"Atari Episode Score: {episode_infos['r']:.2f}")
+                        print("Atari Episode Length", episode_infos["l"])
+
+                if done and not is_atari and args.verbose > 0:
+                    # NOTE: for env using VecNormalize, the mean reward
+                    # is a normalized reward when `--norm_reward` flag is passed
+                    print(f"Episode Reward: {episode_reward:.2f}")
+                    print("Episode Length", ep_len)
+                    episode_rewards.append(episode_reward)
+                    episode_lengths.append(ep_len)
+                    episode_reward = 0.0
+                    ep_len = 0
+
+                # Reset also when the goal is achieved when using HER
+                if done and infos[0].get("is_success") is not None:
+                    if args.verbose > 1:
+                        print("Success?", infos[0].get("is_success", False))
+
+                    if infos[0].get("is_success") is not None:
+                        successes.append(infos[0].get("is_success", False))
+                        episode_reward, ep_len = 0.0, 0
+
+    except KeyboardInterrupt:
+        pass
+
+    if args.verbose > 0 and len(successes) > 0:
+        print(f"Success rate: {100 * np.mean(successes):.2f}%")
+
+    if args.verbose > 0 and len(episode_rewards) > 0:
+        print(f"{len(episode_rewards)} Episodes")
+        print(f"Mean reward: {np.mean(episode_rewards):.2f} +/- {np.std(episode_rewards):.2f}")
+
+    if args.verbose > 0 and len(episode_lengths) > 0:
+        print(f"Mean episode length: {np.mean(episode_lengths):.2f} +/- {np.std(episode_lengths):.2f}")
+
+    env.close()
+
+
+if __name__ == "__main__":
+    enjoy()
diff --git a/utils/exp_manager.py b/rl_zoo/exp_manager.py
similarity index 98%
rename from utils/exp_manager.py
rename to rl_zoo/exp_manager.py
index 08bc97e56..8a661aee4 100644
--- a/utils/exp_manager.py
+++ b/rl_zoo/exp_manager.py
@@ -4,6 +4,7 @@
 import time
 import warnings
 from collections import OrderedDict
+from pathlib import Path
 from pprint import pprint
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -44,10 +45,10 @@
 from torch import nn as nn  # noqa: F401
 
 # Register custom envs
-import utils.import_envs  # noqa: F401 pytype: disable=import-error
-from utils.callbacks import SaveVecNormalizeCallback, TQDMCallback, TrialEvalCallback
-from utils.hyperparams_opt import HYPERPARAMS_SAMPLER
-from utils.utils import ALGOS, get_callback_list, get_latest_run_id, get_wrapper_class, linear_schedule
+import rl_zoo.import_envs  # noqa: F401 pytype: disable=import-error
+from rl_zoo.callbacks import SaveVecNormalizeCallback, TQDMCallback, TrialEvalCallback
+from rl_zoo.hyperparams_opt import HYPERPARAMS_SAMPLER
+from rl_zoo.utils import ALGOS, get_callback_list, get_latest_run_id, get_wrapper_class, linear_schedule
 
 
 class ExperimentManager:
@@ -101,7 +102,8 @@ def __init__(
         self.env_name = EnvironmentName(env_id)
         # Custom params
         self.custom_hyperparams = hyperparams
-        self.yaml_file = yaml_file or f"hyperparams/{self.algo}.yml"
+        default_path = Path(__file__).parent.parent
+        self.yaml_file = yaml_file or str(default_path / f"hyperparams/{self.algo}.yml")
         self.env_kwargs = {} if env_kwargs is None else env_kwargs
         self.n_timesteps = n_timesteps
         self.normalize = False
diff --git a/utils/hyperparams_opt.py b/rl_zoo/hyperparams_opt.py
similarity index 99%
rename from utils/hyperparams_opt.py
rename to rl_zoo/hyperparams_opt.py
index 81add5858..1c7f86796 100644
--- a/utils/hyperparams_opt.py
+++ b/rl_zoo/hyperparams_opt.py
@@ -5,7 +5,7 @@
 from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
 from torch import nn as nn
 
-from utils import linear_schedule
+from rl_zoo import linear_schedule
 
 
 def sample_ppo_params(trial: optuna.Trial) -> Dict[str, Any]:
diff --git a/utils/import_envs.py b/rl_zoo/import_envs.py
similarity index 96%
rename from utils/import_envs.py
rename to rl_zoo/import_envs.py
index 3dc9d7554..8eb1861ae 100644
--- a/utils/import_envs.py
+++ b/rl_zoo/import_envs.py
@@ -1,7 +1,7 @@
 import gym
 from gym.envs.registration import register
 
-from utils.wrappers import MaskVelocityWrapper
+from rl_zoo.wrappers import MaskVelocityWrapper
 
 try:
     import pybullet_envs  # pytype: disable=import-error
diff --git a/utils/load_from_hub.py b/rl_zoo/load_from_hub.py
similarity index 99%
rename from utils/load_from_hub.py
rename to rl_zoo/load_from_hub.py
index a1cc39591..cb1aceba2 100644
--- a/utils/load_from_hub.py
+++ b/rl_zoo/load_from_hub.py
@@ -8,7 +8,7 @@
 from huggingface_sb3 import EnvironmentName, ModelName, ModelRepoId, load_from_hub
 from requests.exceptions import HTTPError
 
-from utils import ALGOS, get_latest_run_id
+from rl_zoo import ALGOS, get_latest_run_id
 
 
 def download_from_hub(
diff --git a/utils/push_to_hub.py b/rl_zoo/push_to_hub.py
similarity index 98%
rename from utils/push_to_hub.py
rename to rl_zoo/push_to_hub.py
index 79fb6f902..9c2ed556a 100644
--- a/utils/push_to_hub.py
+++ b/rl_zoo/push_to_hub.py
@@ -19,10 +19,10 @@
 from stable_baselines3.common.vec_env import VecEnv, unwrap_vec_normalize
 from wasabi import Printer
 
-import utils.import_envs  # noqa: F401 pylint: disable=unused-import
-from utils import ALGOS, create_test_env, get_saved_hyperparams
-from utils.exp_manager import ExperimentManager
-from utils.utils import StoreDict, get_model_path
+import rl_zoo.import_envs  # noqa: F401 pylint: disable=unused-import
+from rl_zoo import ALGOS, create_test_env, get_saved_hyperparams
+from rl_zoo.exp_manager import ExperimentManager
+from rl_zoo.utils import StoreDict, get_model_path
 
 msg = Printer()
 
diff --git a/rl_zoo/py.typed b/rl_zoo/py.typed
new file mode 100644
index 000000000..e69de29bb
diff --git a/utils/record_training.py b/rl_zoo/record_training.py
similarity index 99%
rename from utils/record_training.py
rename to rl_zoo/record_training.py
index f6a3a043f..398cb2064 100644
--- a/utils/record_training.py
+++ b/rl_zoo/record_training.py
@@ -7,7 +7,7 @@
 
 from huggingface_sb3 import EnvironmentName
 
-from utils.utils import ALGOS, get_latest_run_id
+from rl_zoo.utils import ALGOS, get_latest_run_id
 
 if __name__ == "__main__":  # noqa: C901
     parser = argparse.ArgumentParser()
diff --git a/utils/record_video.py b/rl_zoo/record_video.py
similarity index 97%
rename from utils/record_video.py
rename to rl_zoo/record_video.py
index bfa34132d..ea15b1cfb 100644
--- a/utils/record_video.py
+++ b/rl_zoo/record_video.py
@@ -8,8 +8,8 @@
 from stable_baselines3.common.utils import set_random_seed
 from stable_baselines3.common.vec_env import VecVideoRecorder
 
-from utils.exp_manager import ExperimentManager
-from utils.utils import ALGOS, StoreDict, create_test_env, get_model_path, get_saved_hyperparams
+from rl_zoo.exp_manager import ExperimentManager
+from rl_zoo.utils import ALGOS, StoreDict, create_test_env, get_model_path, get_saved_hyperparams
 
 if __name__ == "__main__":  # noqa: C901
     parser = argparse.ArgumentParser()
diff --git a/rl_zoo/train.py b/rl_zoo/train.py
new file mode 100644
index 000000000..79ce938d4
--- /dev/null
+++ b/rl_zoo/train.py
@@ -0,0 +1,255 @@
+import argparse
+import difflib
+import importlib
+import os
+import time
+import uuid
+
+import gym
+import numpy as np
+import torch as th
+from stable_baselines3.common.utils import set_random_seed
+
+# Register custom envs
+import rl_zoo.import_envs  # noqa: F401 pytype: disable=import-error
+from rl_zoo.exp_manager import ExperimentManager
+from rl_zoo.utils import ALGOS, StoreDict
+
+
+def train():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False, choices=list(ALGOS.keys()))
+    parser.add_argument("--env", type=str, default="CartPole-v1", help="environment ID")
+    parser.add_argument("-tb", "--tensorboard-log", help="Tensorboard log dir", default="", type=str)
+    parser.add_argument("-i", "--trained-agent", help="Path to a pretrained agent to continue training", default="", type=str)
+    parser.add_argument(
+        "--truncate-last-trajectory",
+        help="When using HER with online sampling the last trajectory "
+        "in the replay buffer will be truncated after reloading the replay buffer.",
+        default=True,
+        type=bool,
+    )
+    parser.add_argument("-n", "--n-timesteps", help="Overwrite the number of timesteps", default=-1, type=int)
+    parser.add_argument("--num-threads", help="Number of threads for PyTorch (-1 to use default)", default=-1, type=int)
+    parser.add_argument("--log-interval", help="Override log interval (default: -1, no change)", default=-1, type=int)
+    parser.add_argument(
+        "--eval-freq",
+        help="Evaluate the agent every n steps (if negative, no evaluation). "
+        "During hyperparameter optimization n-evaluations is used instead",
+        default=25000,
+        type=int,
+    )
+    parser.add_argument(
+        "--optimization-log-path",
+        help="Path to save the evaluation log and optimal policy for each hyperparameter tried during optimization. "
+        "Disabled if no argument is passed.",
+        type=str,
+    )
+    parser.add_argument("--eval-episodes", help="Number of episodes to use for evaluation", default=5, type=int)
+    parser.add_argument("--n-eval-envs", help="Number of environments for evaluation", default=1, type=int)
+    parser.add_argument("--save-freq", help="Save the model every n steps (if negative, no checkpoint)", default=-1, type=int)
+    parser.add_argument(
+        "--save-replay-buffer", help="Save the replay buffer too (when applicable)", action="store_true", default=False
+    )
+    parser.add_argument("-f", "--log-folder", help="Log folder", type=str, default="logs")
+    parser.add_argument("--seed", help="Random generator seed", type=int, default=-1)
+    parser.add_argument("--vec-env", help="VecEnv type", type=str, default="dummy", choices=["dummy", "subproc"])
+    parser.add_argument("--device", help="PyTorch device to be use (ex: cpu, cuda...)", default="auto", type=str)
+    parser.add_argument(
+        "--n-trials",
+        help="Number of trials for optimizing hyperparameters. "
+        "This applies to each optimization runner, not the entire optimization process.",
+        type=int,
+        default=500,
+    )
+    parser.add_argument(
+        "--max-total-trials",
+        help="Number of (potentially pruned) trials for optimizing hyperparameters. "
+        "This applies to the entire optimization process and takes precedence over --n-trials if set.",
+        type=int,
+        default=None,
+    )
+    parser.add_argument(
+        "-optimize", "--optimize-hyperparameters", action="store_true", default=False, help="Run hyperparameters search"
+    )
+    parser.add_argument(
+        "--no-optim-plots", action="store_true", default=False, help="Disable hyperparameter optimization plots"
+    )
+    parser.add_argument("--n-jobs", help="Number of parallel jobs when optimizing hyperparameters", type=int, default=1)
+    parser.add_argument(
+        "--sampler",
+        help="Sampler to use when optimizing hyperparameters",
+        type=str,
+        default="tpe",
+        choices=["random", "tpe", "skopt"],
+    )
+    parser.add_argument(
+        "--pruner",
+        help="Pruner to use when optimizing hyperparameters",
+        type=str,
+        default="median",
+        choices=["halving", "median", "none"],
+    )
+    parser.add_argument("--n-startup-trials", help="Number of trials before using optuna sampler", type=int, default=10)
+    parser.add_argument(
+        "--n-evaluations",
+        help="Training policies are evaluated every n-timesteps // n-evaluations steps when doing hyperparameter optimization."
+        "Default is 1 evaluation per 100k timesteps.",
+        type=int,
+        default=None,
+    )
+    parser.add_argument(
+        "--storage", help="Database storage path if distributed optimization should be used", type=str, default=None
+    )
+    parser.add_argument("--study-name", help="Study name for distributed optimization", type=str, default=None)
+    parser.add_argument("--verbose", help="Verbose mode (0: no output, 1: INFO)", default=1, type=int)
+    parser.add_argument(
+        "--gym-packages",
+        type=str,
+        nargs="+",
+        default=[],
+        help="Additional external Gym environment package modules to import (e.g. gym_minigrid)",
+    )
+    parser.add_argument(
+        "--env-kwargs", type=str, nargs="+", action=StoreDict, help="Optional keyword argument to pass to the env constructor"
+    )
+    parser.add_argument(
+        "-params",
+        "--hyperparams",
+        type=str,
+        nargs="+",
+        action=StoreDict,
+        help="Overwrite hyperparameter (e.g. learning_rate:0.01 train_freq:10)",
+    )
+    parser.add_argument(
+        "-yaml", "--yaml-file", type=str, default=None, help="Custom yaml file from which the hyperparameters will be loaded"
+    )
+    parser.add_argument("-uuid", "--uuid", action="store_true", default=False, help="Ensure that the run has a unique ID")
+    parser.add_argument(
+        "--track",
+        action="store_true",
+        default=False,
+        help="if toggled, this experiment will be tracked with Weights and Biases",
+    )
+    parser.add_argument("--wandb-project-name", type=str, default="sb3", help="the wandb's project name")
+    parser.add_argument("--wandb-entity", type=str, default=None, help="the entity (team) of wandb's project")
+    parser.add_argument(
+        "-P",
+        "--progress",
+        action="store_true",
+        default=False,
+        help="if toggled, display a progress bar using tqdm and rich",
+    )
+
+    args = parser.parse_args()
+
+    # Going through custom gym packages to let them register in the global registory
+    for env_module in args.gym_packages:
+        importlib.import_module(env_module)
+
+    env_id = args.env
+    registered_envs = set(gym.envs.registry.env_specs.keys())  # pytype: disable=module-attr
+
+    # If the environment is not found, suggest the closest match
+    if env_id not in registered_envs:
+        try:
+            closest_match = difflib.get_close_matches(env_id, registered_envs, n=1)[0]
+        except IndexError:
+            closest_match = "'no close match found...'"
+        raise ValueError(f"{env_id} not found in gym registry, you maybe meant {closest_match}?")
+
+    # Unique id to ensure there is no race condition for the folder creation
+    uuid_str = f"_{uuid.uuid4()}" if args.uuid else ""
+    if args.seed < 0:
+        # Seed but with a random one
+        args.seed = np.random.randint(2**32 - 1, dtype="int64").item()
+
+    set_random_seed(args.seed)
+
+    # Setting num threads to 1 makes things run faster on cpu
+    if args.num_threads > 0:
+        if args.verbose > 1:
+            print(f"Setting torch.num_threads to {args.num_threads}")
+        th.set_num_threads(args.num_threads)
+
+    if args.trained_agent != "":
+        assert args.trained_agent.endswith(".zip") and os.path.isfile(
+            args.trained_agent
+        ), "The trained_agent must be a valid path to a .zip file"
+
+    print("=" * 10, env_id, "=" * 10)
+    print(f"Seed: {args.seed}")
+
+    if args.track:
+        try:
+            import wandb
+        except ImportError:
+            raise ImportError(
+                "if you want to use Weights & Biases to track experiment, please install W&B via `pip install wandb`"
+            )
+
+        run_name = f"{args.env}__{args.algo}__{args.seed}__{int(time.time())}"
+        run = wandb.init(
+            name=run_name,
+            project=args.wandb_project_name,
+            entity=args.wandb_entity,
+            config=vars(args),
+            sync_tensorboard=True,  # auto-upload sb3's tensorboard metrics
+            monitor_gym=True,  # auto-upload the videos of agents playing the game
+            save_code=True,  # optional
+        )
+        args.tensorboard_log = f"runs/{run_name}"
+
+    exp_manager = ExperimentManager(
+        args,
+        args.algo,
+        env_id,
+        args.log_folder,
+        args.tensorboard_log,
+        args.n_timesteps,
+        args.eval_freq,
+        args.eval_episodes,
+        args.save_freq,
+        args.hyperparams,
+        args.env_kwargs,
+        args.trained_agent,
+        args.optimize_hyperparameters,
+        args.storage,
+        args.study_name,
+        args.n_trials,
+        args.max_total_trials,
+        args.n_jobs,
+        args.sampler,
+        args.pruner,
+        args.optimization_log_path,
+        n_startup_trials=args.n_startup_trials,
+        n_evaluations=args.n_evaluations,
+        truncate_last_trajectory=args.truncate_last_trajectory,
+        uuid_str=uuid_str,
+        seed=args.seed,
+        log_interval=args.log_interval,
+        save_replay_buffer=args.save_replay_buffer,
+        verbose=args.verbose,
+        vec_env_type=args.vec_env,
+        n_eval_envs=args.n_eval_envs,
+        no_optim_plots=args.no_optim_plots,
+        device=args.device,
+        yaml_file=args.yaml_file,
+        show_progress=args.progress,
+    )
+
+    # Prepare experiment and launch hyperparameter optimization if needed
+    results = exp_manager.setup_experiment()
+    if results is not None:
+        model, saved_hyperparams = results
+        if args.track:
+            # we need to save the loaded hyperparameters
+            args.saved_hyperparams = saved_hyperparams
+            run.config.setdefaults(vars(args))
+
+        # Normal training
+        if model is not None:
+            exp_manager.learn(model)
+            exp_manager.save_trained_model(model)
+    else:
+        exp_manager.hyperparameters_optimization()
diff --git a/utils/utils.py b/rl_zoo/utils.py
similarity index 99%
rename from utils/utils.py
rename to rl_zoo/utils.py
index 89ad5e332..3da88eabf 100644
--- a/utils/utils.py
+++ b/rl_zoo/utils.py
@@ -199,7 +199,7 @@ def create_test_env(
     :return:
     """
     # Avoid circular import
-    from utils.exp_manager import ExperimentManager
+    from rl_zoo.exp_manager import ExperimentManager
 
     # Create the environment and wrap it if necessary
     env_wrapper = get_wrapper_class(hyperparams)
diff --git a/rl_zoo/version.txt b/rl_zoo/version.txt
new file mode 100644
index 000000000..fdd3be6df
--- /dev/null
+++ b/rl_zoo/version.txt
@@ -0,0 +1,1 @@
+1.6.2
diff --git a/utils/wrappers.py b/rl_zoo/wrappers.py
similarity index 83%
rename from utils/wrappers.py
rename to rl_zoo/wrappers.py
index 8f1e6c5a8..960020dc2 100644
--- a/utils/wrappers.py
+++ b/rl_zoo/wrappers.py
@@ -1,7 +1,6 @@
 import gym
 import numpy as np
 from sb3_contrib.common.wrappers import TimeFeatureWrapper  # noqa: F401 (backward compatibility)
-from scipy.signal import iirfilter, sosfilt, zpk2sos
 
 
 class DoneOnSuccessWrapper(gym.Wrapper):
@@ -55,70 +54,6 @@ def step(self, action):
         return self.env.step(noisy_action)
 
 
-# from https://docs.obspy.org
-def lowpass(data, freq, df, corners=4, zerophase=False):
-    """
-    Butterworth-Lowpass Filter.
-
-    Filter data removing data over certain frequency ``freq`` using ``corners``
-    corners.
-    The filter uses :func:`scipy.signal.iirfilter` (for design)
-    and :func:`scipy.signal.sosfilt` (for applying the filter).
-
-    :type data: numpy.ndarray
-    :param data: Data to filter.
-    :param freq: Filter corner frequency.
-    :param df: Sampling rate in Hz.
-    :param corners: Filter corners / order.
-    :param zerophase: If True, apply filter once forwards and once backwards.
-        This results in twice the number of corners but zero phase shift in
-        the resulting filtered trace.
-    :return: Filtered data.
-    """
-    fe = 0.5 * df
-    f = freq / fe
-    # raise for some bad scenarios
-    if f > 1:
-        f = 1.0
-        msg = "Selected corner frequency is above Nyquist. " + "Setting Nyquist as high corner."
-        print(msg)
-    z, p, k = iirfilter(corners, f, btype="lowpass", ftype="butter", output="zpk")
-    sos = zpk2sos(z, p, k)
-    if zerophase:
-        firstpass = sosfilt(sos, data)
-        return sosfilt(sos, firstpass[::-1])[::-1]
-    else:
-        return sosfilt(sos, data)
-
-
-class LowPassFilterWrapper(gym.Wrapper):
-    """
-    Butterworth-Lowpass
-
-    :param env:
-    :param freq: Filter corner frequency.
-    :param df: Sampling rate in Hz.
-    """
-
-    def __init__(self, env: gym.Env, freq: float = 5.0, df: float = 25.0):
-        super().__init__(env)
-        self.freq = freq
-        self.df = df
-        self.signal = []
-
-    def reset(self):
-        self.signal = []
-        return self.env.reset()
-
-    def step(self, action):
-        self.signal.append(action)
-        filtered = np.zeros_like(action)
-        for i in range(self.action_space.shape[0]):
-            smoothed_action = lowpass(np.array(self.signal)[:, i], freq=self.freq, df=self.df)
-            filtered[i] = smoothed_action[-1]
-        return self.env.step(filtered)
-
-
 class ActionSmoothingWrapper(gym.Wrapper):
     """
     Smooth the action using exponential moving average.
diff --git a/scripts/migrate_to_hub.py b/scripts/migrate_to_hub.py
index 81bef4d6d..921952c1b 100644
--- a/scripts/migrate_to_hub.py
+++ b/scripts/migrate_to_hub.py
@@ -1,6 +1,6 @@
 import subprocess
 
-from utils.utils import get_hf_trained_models, get_trained_models
+from rl_zoo.utils import get_hf_trained_models, get_trained_models
 
 folder = "rl-trained-agents"
 orga = "sb3"
diff --git a/setup.cfg b/setup.cfg
index c74b8320e..b2796ea5e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -15,8 +15,8 @@ inputs = .
 ignore = W503,W504,E203,E231  # line breaks before and after binary operators
 # Ignore import not used when aliases are defined
 per-file-ignores =
-	./utils/__init__.py:F401
-  ./utils/import_envs.py:F401
+	./rl_zoo/__init__.py:F401
+  ./rl_zoo/import_envs.py:F401
   ./scripts/all_plots.py:E501
   ./scripts/plot_train.py:E501
   ./scripts/plot_training_success.py:E501
@@ -33,4 +33,4 @@ max-line-length = 127
 [isort]
 profile = black
 line_length = 127
-src_paths = stable_baselines3,utils
+src_paths = stable_baselines3,rl_zoo
diff --git a/setup.py b/setup.py
new file mode 100644
index 000000000..3787af35a
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,65 @@
+import os
+
+from setuptools import find_packages, setup
+
+with open(os.path.join("rl_zoo", "version.txt")) as file_handler:
+    __version__ = file_handler.read().strip()
+
+
+long_description = """
+# RL Baselines3 Zoo: A Training Framework for Stable Baselines3 Reinforcement Learning Agents
+
+See https://github.com/DLR-RM/rl-baselines3-zoo
+"""
+
+setup(
+    name="rl_zoo",
+    packages=[package for package in find_packages() if package.startswith("rl_zoo")],
+    package_data={
+        "rl_zoo": [
+            "py.typed",
+            "version.txt",
+            "../scripts/*.py",
+            "../hyperparams/*.yml",
+        ]
+    },
+    scripts=[
+        "./scripts/all_plots.py",
+        "./scripts/plot_train.py",
+        "./scripts/plot_from_file.py",
+    ],
+    entry_points={"console_scripts": ["rl_zoo_train=rl_zoo.train:train", "rl_zoo=rl_zoo.cli:main"]},
+    install_requires=[
+        "sb3-contrib>=1.6.1",
+        "huggingface_sb3>=2.2.1, <3.*",
+        "tqdm",
+        "rich",
+        "optuna",
+        "pyyaml>=5.1",
+        "pytablewriter~=0.64",
+        # TODO: add test dependencies
+    ],
+    description="A Training Framework for Stable Baselines3 Reinforcement Learning Agents",
+    author="Antonin Raffin",
+    url="https://github.com/DLR-RM/rl-baselines3-zoo",
+    author_email="antonin.raffin@dlr.de",
+    keywords="reinforcement-learning-algorithms reinforcement-learning machine-learning "
+    "gym openai stable baselines toolbox python data-science",
+    license="MIT",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    version=__version__,
+    python_requires=">=3.7",
+    # PyPI package information.
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+    ],
+)
+
+# python setup.py sdist
+# python setup.py bdist_wheel
+# twine upload --repository-url https://test.pypi.org/legacy/ dist/*
+# twine upload dist/*
diff --git a/train.py b/train.py
index bf6b0a6cc..7316e13e1 100644
--- a/train.py
+++ b/train.py
@@ -1,257 +1,4 @@
-import argparse
-import difflib
-import importlib
-import os
-import time
-import uuid
-
-import gym
-import numpy as np
-import seaborn
-import torch as th
-from stable_baselines3.common.utils import set_random_seed
-
-# Register custom envs
-import utils.import_envs  # noqa: F401 pytype: disable=import-error
-from utils.exp_manager import ExperimentManager
-from utils.utils import ALGOS, StoreDict
-
-seaborn.set()
+from rl_zoo.train import train
 
 if __name__ == "__main__":  # noqa: C901
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False, choices=list(ALGOS.keys()))
-    parser.add_argument("--env", type=str, default="CartPole-v1", help="environment ID")
-    parser.add_argument("-tb", "--tensorboard-log", help="Tensorboard log dir", default="", type=str)
-    parser.add_argument("-i", "--trained-agent", help="Path to a pretrained agent to continue training", default="", type=str)
-    parser.add_argument(
-        "--truncate-last-trajectory",
-        help="When using HER with online sampling the last trajectory "
-        "in the replay buffer will be truncated after reloading the replay buffer.",
-        default=True,
-        type=bool,
-    )
-    parser.add_argument("-n", "--n-timesteps", help="Overwrite the number of timesteps", default=-1, type=int)
-    parser.add_argument("--num-threads", help="Number of threads for PyTorch (-1 to use default)", default=-1, type=int)
-    parser.add_argument("--log-interval", help="Override log interval (default: -1, no change)", default=-1, type=int)
-    parser.add_argument(
-        "--eval-freq",
-        help="Evaluate the agent every n steps (if negative, no evaluation). "
-        "During hyperparameter optimization n-evaluations is used instead",
-        default=25000,
-        type=int,
-    )
-    parser.add_argument(
-        "--optimization-log-path",
-        help="Path to save the evaluation log and optimal policy for each hyperparameter tried during optimization. "
-        "Disabled if no argument is passed.",
-        type=str,
-    )
-    parser.add_argument("--eval-episodes", help="Number of episodes to use for evaluation", default=5, type=int)
-    parser.add_argument("--n-eval-envs", help="Number of environments for evaluation", default=1, type=int)
-    parser.add_argument("--save-freq", help="Save the model every n steps (if negative, no checkpoint)", default=-1, type=int)
-    parser.add_argument(
-        "--save-replay-buffer", help="Save the replay buffer too (when applicable)", action="store_true", default=False
-    )
-    parser.add_argument("-f", "--log-folder", help="Log folder", type=str, default="logs")
-    parser.add_argument("--seed", help="Random generator seed", type=int, default=-1)
-    parser.add_argument("--vec-env", help="VecEnv type", type=str, default="dummy", choices=["dummy", "subproc"])
-    parser.add_argument("--device", help="PyTorch device to be use (ex: cpu, cuda...)", default="auto", type=str)
-    parser.add_argument(
-        "--n-trials",
-        help="Number of trials for optimizing hyperparameters. "
-        "This applies to each optimization runner, not the entire optimization process.",
-        type=int,
-        default=500,
-    )
-    parser.add_argument(
-        "--max-total-trials",
-        help="Number of (potentially pruned) trials for optimizing hyperparameters. "
-        "This applies to the entire optimization process and takes precedence over --n-trials if set.",
-        type=int,
-        default=None,
-    )
-    parser.add_argument(
-        "-optimize", "--optimize-hyperparameters", action="store_true", default=False, help="Run hyperparameters search"
-    )
-    parser.add_argument(
-        "--no-optim-plots", action="store_true", default=False, help="Disable hyperparameter optimization plots"
-    )
-    parser.add_argument("--n-jobs", help="Number of parallel jobs when optimizing hyperparameters", type=int, default=1)
-    parser.add_argument(
-        "--sampler",
-        help="Sampler to use when optimizing hyperparameters",
-        type=str,
-        default="tpe",
-        choices=["random", "tpe", "skopt"],
-    )
-    parser.add_argument(
-        "--pruner",
-        help="Pruner to use when optimizing hyperparameters",
-        type=str,
-        default="median",
-        choices=["halving", "median", "none"],
-    )
-    parser.add_argument("--n-startup-trials", help="Number of trials before using optuna sampler", type=int, default=10)
-    parser.add_argument(
-        "--n-evaluations",
-        help="Training policies are evaluated every n-timesteps // n-evaluations steps when doing hyperparameter optimization."
-        "Default is 1 evaluation per 100k timesteps.",
-        type=int,
-        default=None,
-    )
-    parser.add_argument(
-        "--storage", help="Database storage path if distributed optimization should be used", type=str, default=None
-    )
-    parser.add_argument("--study-name", help="Study name for distributed optimization", type=str, default=None)
-    parser.add_argument("--verbose", help="Verbose mode (0: no output, 1: INFO)", default=1, type=int)
-    parser.add_argument(
-        "--gym-packages",
-        type=str,
-        nargs="+",
-        default=[],
-        help="Additional external Gym environment package modules to import (e.g. gym_minigrid)",
-    )
-    parser.add_argument(
-        "--env-kwargs", type=str, nargs="+", action=StoreDict, help="Optional keyword argument to pass to the env constructor"
-    )
-    parser.add_argument(
-        "-params",
-        "--hyperparams",
-        type=str,
-        nargs="+",
-        action=StoreDict,
-        help="Overwrite hyperparameter (e.g. learning_rate:0.01 train_freq:10)",
-    )
-    parser.add_argument(
-        "-yaml", "--yaml-file", type=str, default=None, help="Custom yaml file from which the hyperparameters will be loaded"
-    )
-    parser.add_argument("-uuid", "--uuid", action="store_true", default=False, help="Ensure that the run has a unique ID")
-    parser.add_argument(
-        "--track",
-        action="store_true",
-        default=False,
-        help="if toggled, this experiment will be tracked with Weights and Biases",
-    )
-    parser.add_argument("--wandb-project-name", type=str, default="sb3", help="the wandb's project name")
-    parser.add_argument("--wandb-entity", type=str, default=None, help="the entity (team) of wandb's project")
-    parser.add_argument(
-        "-P",
-        "--progress",
-        action="store_true",
-        default=False,
-        help="if toggled, display a progress bar using tqdm and rich",
-    )
-
-    args = parser.parse_args()
-
-    # Going through custom gym packages to let them register in the global registory
-    for env_module in args.gym_packages:
-        importlib.import_module(env_module)
-
-    env_id = args.env
-    registered_envs = set(gym.envs.registry.env_specs.keys())  # pytype: disable=module-attr
-
-    # If the environment is not found, suggest the closest match
-    if env_id not in registered_envs:
-        try:
-            closest_match = difflib.get_close_matches(env_id, registered_envs, n=1)[0]
-        except IndexError:
-            closest_match = "'no close match found...'"
-        raise ValueError(f"{env_id} not found in gym registry, you maybe meant {closest_match}?")
-
-    # Unique id to ensure there is no race condition for the folder creation
-    uuid_str = f"_{uuid.uuid4()}" if args.uuid else ""
-    if args.seed < 0:
-        # Seed but with a random one
-        args.seed = np.random.randint(2**32 - 1, dtype="int64").item()
-
-    set_random_seed(args.seed)
-
-    # Setting num threads to 1 makes things run faster on cpu
-    if args.num_threads > 0:
-        if args.verbose > 1:
-            print(f"Setting torch.num_threads to {args.num_threads}")
-        th.set_num_threads(args.num_threads)
-
-    if args.trained_agent != "":
-        assert args.trained_agent.endswith(".zip") and os.path.isfile(
-            args.trained_agent
-        ), "The trained_agent must be a valid path to a .zip file"
-
-    print("=" * 10, env_id, "=" * 10)
-    print(f"Seed: {args.seed}")
-
-    if args.track:
-        try:
-            import wandb
-        except ImportError:
-            raise ImportError(
-                "if you want to use Weights & Biases to track experiment, please install W&B via `pip install wandb`"
-            )
-
-        run_name = f"{args.env}__{args.algo}__{args.seed}__{int(time.time())}"
-        run = wandb.init(
-            name=run_name,
-            project=args.wandb_project_name,
-            entity=args.wandb_entity,
-            config=vars(args),
-            sync_tensorboard=True,  # auto-upload sb3's tensorboard metrics
-            monitor_gym=True,  # auto-upload the videos of agents playing the game
-            save_code=True,  # optional
-        )
-        args.tensorboard_log = f"runs/{run_name}"
-
-    exp_manager = ExperimentManager(
-        args,
-        args.algo,
-        env_id,
-        args.log_folder,
-        args.tensorboard_log,
-        args.n_timesteps,
-        args.eval_freq,
-        args.eval_episodes,
-        args.save_freq,
-        args.hyperparams,
-        args.env_kwargs,
-        args.trained_agent,
-        args.optimize_hyperparameters,
-        args.storage,
-        args.study_name,
-        args.n_trials,
-        args.max_total_trials,
-        args.n_jobs,
-        args.sampler,
-        args.pruner,
-        args.optimization_log_path,
-        n_startup_trials=args.n_startup_trials,
-        n_evaluations=args.n_evaluations,
-        truncate_last_trajectory=args.truncate_last_trajectory,
-        uuid_str=uuid_str,
-        seed=args.seed,
-        log_interval=args.log_interval,
-        save_replay_buffer=args.save_replay_buffer,
-        verbose=args.verbose,
-        vec_env_type=args.vec_env,
-        n_eval_envs=args.n_eval_envs,
-        no_optim_plots=args.no_optim_plots,
-        device=args.device,
-        yaml_file=args.yaml_file,
-        show_progress=args.progress,
-    )
-
-    # Prepare experiment and launch hyperparameter optimization if needed
-    results = exp_manager.setup_experiment()
-    if results is not None:
-        model, saved_hyperparams = results
-        if args.track:
-            # we need to save the loaded hyperparameters
-            args.saved_hyperparams = saved_hyperparams
-            run.config.setdefaults(vars(args))
-
-        # Normal training
-        if model is not None:
-            exp_manager.learn(model)
-            exp_manager.save_trained_model(model)
-    else:
-        exp_manager.hyperparameters_optimization()
+    train()
diff --git a/utils/__init__.py b/utils/__init__.py
deleted file mode 100644
index b55f8bc36..000000000
--- a/utils/__init__.py
+++ /dev/null
@@ -1,9 +0,0 @@
-from .utils import (
-    ALGOS,
-    create_test_env,
-    get_latest_run_id,
-    get_saved_hyperparams,
-    get_trained_models,
-    get_wrapper_class,
-    linear_schedule,
-)
diff --git a/version.txt b/version.txt
deleted file mode 100644
index 9c6d6293b..000000000
--- a/version.txt
+++ /dev/null
@@ -1,1 +0,0 @@
-1.6.1


Test patch:
diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py
index 737280194..08f8f316b 100644
--- a/tests/test_callbacks.py
+++ b/tests/test_callbacks.py
@@ -14,7 +14,7 @@ def test_raw_stat_callback(tmp_path):
         "--env",
         "CartPole-v1",
         "-params",
-        "callback:'utils.callbacks.RawStatisticsCallback'",
+        "callback:'rl_zoo.callbacks.RawStatisticsCallback'",
         "--tensorboard-log",
         f"{tmp_path}",
     ]
@@ -32,7 +32,7 @@ def test_tqdm_callback(tmp_path):
         "--env",
         "CartPole-v1",
         "-params",
-        "callback:'utils.callbacks.TQDMCallback'",
+        "callback:'rl_zoo.callbacks.TQDMCallback'",
         "--tensorboard-log",
         f"{tmp_path}",
     ]
diff --git a/tests/test_enjoy.py b/tests/test_enjoy.py
index f555232d6..ab0aec2c5 100644
--- a/tests/test_enjoy.py
+++ b/tests/test_enjoy.py
@@ -3,7 +3,7 @@
 
 import pytest
 
-from utils.utils import get_hf_trained_models, get_trained_models
+from rl_zoo.utils import get_hf_trained_models, get_trained_models
 
 
 def _assert_eq(left, right):
@@ -46,7 +46,7 @@ def test_trained_agents(trained_model):
 def test_benchmark(tmp_path):
     args = ["-n", str(N_STEPS), "--benchmark-dir", tmp_path, "--test-mode", "--no-hub"]
 
-    return_code = subprocess.call(["python", "-m", "utils.benchmark"] + args)
+    return_code = subprocess.call(["python", "-m", "rl_zoo.benchmark"] + args)
     _assert_eq(return_code, 0)
 
 
@@ -94,7 +94,7 @@ def test_record_video(tmp_path):
     # Skip if no X-Server
     pytest.importorskip("pyglet.gl")
 
-    return_code = subprocess.call(["python", "-m", "utils.record_video"] + args)
+    return_code = subprocess.call(["python", "-m", "rl_zoo.record_video"] + args)
     _assert_eq(return_code, 0)
     video_path = str(tmp_path / "final-model-sac-Pendulum-v1-step-0-to-step-100.mp4")
     # File is not empty
@@ -135,7 +135,7 @@ def test_record_training(tmp_path):
     return_code = subprocess.call(["python", "train.py"] + args_training)
     _assert_eq(return_code, 0)
 
-    return_code = subprocess.call(["python", "-m", "utils.record_training"] + args_recording)
+    return_code = subprocess.call(["python", "-m", "rl_zoo.record_training"] + args_recording)
     _assert_eq(return_code, 0)
     mp4_path = str(videos_tmp_path / "training.mp4")
     gif_path = str(videos_tmp_path / "training.gif")
diff --git a/tests/test_train.py b/tests/test_train.py
index 4973565d2..ed32d144c 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -96,7 +96,7 @@ def test_parallel_train(tmp_path):
         "--log-folder",
         tmp_path,
         "-params",
-        "callback:'utils.callbacks.ParallelTrainCallback'",
+        "callback:'rl_zoo.callbacks.ParallelTrainCallback'",
     ]
 
     return_code = subprocess.call(["python", "train.py"] + args)
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
index bfe0d7a1b..929c75592 100644
--- a/tests/test_wrappers.py
+++ b/tests/test_wrappers.py
@@ -5,8 +5,8 @@
 from stable_baselines3.common.env_checker import check_env
 from stable_baselines3.common.env_util import DummyVecEnv
 
-from utils.utils import get_wrapper_class
-from utils.wrappers import ActionNoiseWrapper, DelayedRewardWrapper, HistoryWrapper, TimeFeatureWrapper
+from rl_zoo.utils import get_wrapper_class
+from rl_zoo.wrappers import ActionNoiseWrapper, DelayedRewardWrapper, HistoryWrapper, TimeFeatureWrapper
 
 
 def test_wrappers():
@@ -22,8 +22,8 @@ def test_wrappers():
     "env_wrapper",
     [
         None,
-        {"utils.wrappers.HistoryWrapper": dict(horizon=2)},
-        [{"utils.wrappers.HistoryWrapper": dict(horizon=3)}, "utils.wrappers.TimeFeatureWrapper"],
+        {"rl_zoo.wrappers.HistoryWrapper": dict(horizon=2)},
+        [{"rl_zoo.wrappers.HistoryWrapper": dict(horizon=3)}, "rl_zoo.wrappers.TimeFeatureWrapper"],
     ],
 )
 def test_get_wrapper(env_wrapper):