Automatic download of TmrlData when missing

Author: yannbouteiller

readme/Install.md

@@ -1,7 +1,7 @@
 # Prerequisites
 * Windows / Linux
 * Python >= 3.7
-* A recent NVIDIA GPU (required only on the training computer if you plan to train your own AI)
+* A recent NVIDIA GPU (required only on the training computer if you plan to train your own models)
 
 #### If using Anaconda on Windows:
 
@@ -47,6 +47,13 @@ To install the `tmrl` python library, open your favorite terminal and run:
 ```shell
 pip install tmrl
 ```
+
+Then, validate the installation:
+
+```shell
+python -m tmrl --install
+```
+
 #### Additional information for Windows / Trackmania 2020:
 
 If running on Windows, during the installation, a driver will be installed to emulate a virtual gamepad.
@@ -74,6 +81,8 @@ If at some point you want to do a clean re-install of `tmrl`:
 - Delete the `TmrlData` folder from your home folder
 - `pip install tmrl`
 
+## Set up TMRL
+
 ### (Optional) Configure/manage TMRL:
 
 The `TmrlData` folder is your _"control pannel"_, it contains everything `tmrl` uses and generates:
@@ -95,6 +104,15 @@ In particular, you may want to adapt the following entries:
 
 You can delete the content of all folders (but not the folders themselves) whenever you like (except `config.json`, a default version is provided in `resources` if you delete this).
 
+To reset the library, delete the entire `TmrlData` folder and run:
+
+```shell
+python -m tmrl --install
+```
+
+This will download and extract the `TmrlData` folder back to its original state.
+
+
 ### (Optional) Check that everything works:
 
 Launch TrackMania 2020, launch a track, then press `f3` to open the OpenPlanet menu, open the logs by clicking `OpenPlanet > Log`, and in the OpenPlanet menu click `Developer > (Re)load plugin > TMRL Grab Data`.

readme/install_linux.md

@@ -74,7 +74,13 @@ Open a terminal and run:
 ```bash
 pip3 install tmrl
 ```
-This installs the `tmrl` library in your active python environment and creates a `TmrlData` folder in your home directory.
+
+Validate:
+```bash
+python -m tmrl --install
+```
+
+The `tmrl` library is now installed in your active python environment and has created a `TmrlData` folder in your home directory.
 
 Navigate to your TrackMania Proton folder:
 
@@ -94,30 +100,4 @@ cp ~/TmrlData/resources/tmrl-test.Map.Gbx Documents/Trackmania/Maps/My Maps/.
 
 ## Set up `tmrl`
 
-
-### Configure/manage TMRL:
-
-The `TmrlData` folder is your _"control pannel"_, it contains everything `tmrl` uses and generates:
-- The `checkpoints` subfolder is used by the trainer process: it contains persistent checkpoints of your training,
-- The `weights` subfolder is used by the worker process: it contains snapshots of your trained policies,
-- The `reward` subfolder is used by the worker process: it contains your reward function,
-- The `dataset` subfolder is for RL developers (to use with custom replay buffers),
-- The `config` subfolder contains a configuration file that you probably want to tweak.
-
-Navigate to `TmrlData\config` and open `config.json` in a text editor.
-
-In particular, you may want to adapt the following entries:
-- `RUN_NAME`: set a new name for starting training from scratch
-- `LOCALHOST_WORKER`: set to `false` for `workers` not on the same computer as the `server`
-- `LOCALHOST_TRAINER`: set to `false` for `trainer` not on the same computer as the `server`
-- `PUBLIC_IP_SERVER`: public IP of the `server` if not running on localhost
-- `PORT` needs to be forwarded on the `server` if not running on localhost
-- `WANDB_PROJECT`, `WANDB_ENTITY` and `WANDB_KEY` can be replaced by you own [wandb](https://wandb.ai/site) credentials for monitoring training
-
-You can delete the content of all folders (but not the folders themselves) whenever you like (except `config.json`, a default version is provided in `resources` if you delete this).
-
-### Check that everything works:
-
-Launch TrackMania 2020, launch a track, then press `f3` to open the OpenPlanet menu, open the logs by clicking `OpenPlanet > Log`, and in the OpenPlanet menu click `Developer > (Re)load plugin > TMRL Grab Data`.
-You should see a message like "waiting for incoming connection" appear in the logs.
-Press `f3` again to close the menu.
+Find out how to configure the library [here](Install.md#set-up-tmrl).

setup.py

@@ -13,6 +13,10 @@
     sys.exit('Sorry, Python < 3.7 is not supported.')
 
 
+# NB: the following code is duplicated under tmrl.tools.init_package.init_tmrl,
+# don't forget to update both whenever changing RESOURCES_URL.
+
+
 RESOURCES_URL = "https://github.com/trackmania-rl/tmrl/releases/download/v0.6.0/resources.zip"
 
 
@@ -133,13 +137,13 @@ def url_retrieve(url: str, outfile: Path, overwrite: bool = False):
 
 setup(
     name='tmrl',
-    version='0.6.1',
+    version='0.6.2',
     description='Network-based framework for real-time robot learning',
     long_description=README,
     long_description_content_type='text/markdown',
     keywords='reinforcement learning, robot learning, trackmania, self driving, roborace',
     url='https://github.com/trackmania-rl/tmrl',
-    download_url='https://github.com/trackmania-rl/tmrl/archive/refs/tags/v0.6.1.tar.gz',
+    download_url='https://github.com/trackmania-rl/tmrl/archive/refs/tags/v0.6.2.tar.gz',
     author='Yann Bouteiller, Edouard Geze',
     author_email='[email protected], [email protected]',
     license='MIT',

tmrl/__init__.py

@@ -25,6 +25,8 @@
             \nPlease install pywin32 manually.")
             raise RuntimeError("Please install pywin32 manually: https://github.com/mhammond/pywin32")
 
+# TMRL folder initialization:
+from tmrl.tools.init_package.init_tmrl import TMRL_FOLDER
 
 # do not remove this
 from dataclasses import dataclass

tmrl/__main__.py

@@ -63,12 +63,15 @@ def main(args):
             check_env_tm20lidar()
         else:
             check_env_tm20full()
+    elif args.install:
+        logging.info(f"TMRL folder: {cfg.TMRL_FOLDER}")
     else:
         raise ArgumentTypeError('Enter a valid argument')
 
 
 if __name__ == "__main__":
     parser = ArgumentParser()
+    parser.add_argument('--install', action='store_true', help='checks TMRL installation')
     parser.add_argument('--server', action='store_true', help='launches the server')
     parser.add_argument('--trainer', action='store_true', help='launches the trainer')
     parser.add_argument('--worker', action='store_true', help='launches a rollout worker')

tmrl/config/config_constants.py

@@ -1,18 +1,26 @@
 # standard library imports
+
+import logging
+
 import os
 from pathlib import Path
-import logging
 import json
 import platform
 from packaging import version
 
 
 __compatibility__ = "0.6.0"
 
+# TMRL FOLDER: =======================================================
+
 SYSTEM = platform.system()
 RTGYM_VERSION = "real-time-gym-v1" if SYSTEM == "Windows" else "real-time-gym-ts-v1"
 
 TMRL_FOLDER = Path.home() / "TmrlData"
+
+if not TMRL_FOLDER.exists():
+    raise RuntimeError(f"Missing folder: {TMRL_FOLDER}")
+
 CHECKPOINTS_FOLDER = TMRL_FOLDER / "checkpoints"
 DATASET_FOLDER = TMRL_FOLDER / "dataset"
 REWARD_FOLDER = TMRL_FOLDER / "reward"

tmrl/tools/init_package/init_pywin32.py

@@ -1,6 +1,6 @@
 # Adapted from https://github.com/mhammond/pywin32/blob/main/pywin32_postinstall.py
 
-# May have to be adapted in the future if becomming incompatible with new version of pywin32.
+# May have to be adapted in the future if becoming incompatible with new versions of pywin32.
 
 # postinstall script for pywin32
 #

tmrl/tools/init_package/init_tmrl.py

@@ -0,0 +1,122 @@
+import logging
+
+import platform
+from pathlib import Path
+
+
+def rmdir(directory):
+    directory = Path(directory)
+    for item in directory.iterdir():
+        if item.is_dir():
+            rmdir(item)
+        else:
+            item.unlink()
+    directory.rmdir()
+
+
+def init_tmrl_data():
+    """
+    Wipes and re-generates the TmrlData folder.
+    """
+    from shutil import copy2
+    from zipfile import ZipFile
+    import urllib.request
+    import urllib.error
+    import socket
+
+    resources_url = "https://github.com/trackmania-rl/tmrl/releases/download/v0.6.0/resources.zip"
+
+    def url_retrieve(url: str, outfile: Path, overwrite: bool = False):
+        """
+        Adapted from https://www.scivision.dev/python-switch-urlretrieve-requests-timeout/
+        """
+        outfile = Path(outfile).expanduser().resolve()
+        if outfile.is_dir():
+            raise ValueError("Please specify full filepath, including filename")
+        if overwrite or not outfile.is_file():
+            outfile.parent.mkdir(parents=True, exist_ok=True)
+            try:
+                urllib.request.urlretrieve(url, str(outfile))
+            except (socket.gaierror, urllib.error.URLError) as err:
+                raise ConnectionError(f"could not download {url} due to {err}")
+
+    # destination folder:
+    home_folder = Path.home()
+    tmrl_folder = home_folder / "TmrlData"
+
+    # Wipe the tmrl folder:
+    if tmrl_folder.exists():
+        rmdir(tmrl_folder)
+
+    # download relevant items IF THE tmrl FOLDER DOESN'T EXIST:
+    assert not tmrl_folder.exists(), f"Failed to delete {tmrl_folder}"
+
+    checkpoints_folder = tmrl_folder / "checkpoints"
+    dataset_folder = tmrl_folder / "dataset"
+    reward_folder = tmrl_folder / "reward"
+    weights_folder = tmrl_folder / "weights"
+    config_folder = tmrl_folder / "config"
+    checkpoints_folder.mkdir(parents=True, exist_ok=True)
+    dataset_folder.mkdir(parents=True, exist_ok=True)
+    reward_folder.mkdir(parents=True, exist_ok=True)
+    weights_folder.mkdir(parents=True, exist_ok=True)
+    config_folder.mkdir(parents=True, exist_ok=True)
+
+    # download resources:
+    resources_target = tmrl_folder / "resources.zip"
+    url_retrieve(resources_url, resources_target)
+
+    # unzip downloaded resources:
+    with ZipFile(resources_target, 'r') as zip_ref:
+        zip_ref.extractall(tmrl_folder)
+
+    # delete zip file:
+    resources_target.unlink()
+
+    # copy relevant files:
+    resources_folder = tmrl_folder / "resources"
+    copy2(resources_folder / "config.json", config_folder)
+    copy2(resources_folder / "reward.pkl", reward_folder)
+    copy2(resources_folder / "SAC_4_LIDAR_pretrained.tmod", weights_folder)
+    copy2(resources_folder / "SAC_4_imgs_pretrained.tmod", weights_folder)
+
+    # on Windows, look for OpenPlanet:
+    if platform.system() == "Windows":
+        openplanet_folder = home_folder / "OpenplanetNext"
+
+        if openplanet_folder.exists():
+            # copy the OpenPlanet script:
+            try:
+                # remove old script if found
+                op_scripts_folder = openplanet_folder / 'Scripts'
+                if op_scripts_folder.exists():
+                    to_remove = [op_scripts_folder / 'Plugin_GrabData_0_1.as',
+                                 op_scripts_folder / 'Plugin_GrabData_0_1.as.sig',
+                                 op_scripts_folder / 'Plugin_GrabData_0_2.as',
+                                 op_scripts_folder / 'Plugin_GrabData_0_2.as.sig']
+                    for old_file in to_remove:
+                        if old_file.exists():
+                            old_file.unlink()
+                # copy new plugin
+                op_plugins_folder = openplanet_folder / 'Plugins'
+                op_plugins_folder.mkdir(parents=True, exist_ok=True)
+                tm20_plugin_1 = resources_folder / 'Plugins' / 'TMRL_GrabData.op'
+                tm20_plugin_2 = resources_folder / 'Plugins' / 'TMRL_SaveGhost.op'
+                copy2(tm20_plugin_1, op_plugins_folder)
+                copy2(tm20_plugin_2, op_plugins_folder)
+            except Exception as e:
+                print(
+                    f"An exception was caught when trying to copy the OpenPlanet plugin automatically. \
+                    Please copy the plugin manually for TrackMania 2020 support. The caught exception was: {str(e)}.")
+        else:
+            # warn the user that OpenPlanet couldn't be found:
+            print(f"The OpenPlanet folder was not found at {openplanet_folder}. \
+            Please copy the OpenPlanet script and signature manually for TrackMania 2020 support.")
+
+
+TMRL_FOLDER = Path.home() / "TmrlData"
+
+if not TMRL_FOLDER.exists():
+    logging.warning(f"The TMRL folder was not found on your machine. Attempting download...")
+    init_tmrl_data()
+    logging.info(f"TMRL folder successfully downloaded, please wait for initialization to complete...")