[WIP] Unity Environment Registry

ml-agents-envs/mlagents_envs/registry/__init__.py

@@ -0,0 +1,4 @@
+from mlagents_envs.registry.unity_env_registry import (  # noqa F401
+    default_registry,
+    UnityEnvRegistry,
+)

ml-agents-envs/mlagents_envs/registry/base_registry_entry.py

@@ -0,0 +1,51 @@
+from abc import abstractmethod
+from typing import Any
+from mlagents_envs.base_env import BaseEnv
+
+
+class BaseRegistryEntry:
+    def __init__(self, identifier: str, expected_reward: float, description: str):
+        """
+        BaseRegistryEntry allows launching a Unity Environment with its make method.
+        :param identifier: The name of the Unity Environment.
+        :param expected_reward: The cumulative reward that an Agent must receive
+        for the task to be considered solved.
+        :param description: A description of the Unity Environment. Contains human
+        readable information about potential special arguments that the make method can
+        take as well as information regarding the observation, reward, actions,
+        behaviors and number of agents in the Environment.
+        """
+        self._identifier = identifier
+        self._expected_reward = expected_reward
+        self._description = description
+
+    @property
+    def identifier(self) -> str:
+        """
+        The unique identifier of the entry
+        """
+        return self._identifier
+
+    @property
+    def expected_reward(self) -> float:
+        """
+        The cumulative reward that an Agent must receive for the task to be considered
+        solved.
+        """
+        return self._expected_reward
+
+    @property
+    def description(self) -> str:
+        """
+        A description of the Unity Environment the entry can make.
+        """
+        return self._description
+
+    @abstractmethod
+    def make(self, **kwargs: Any) -> BaseEnv:
+        """
+        This method creates a Unity BaseEnv (usually a UnityEnvironment).
+        """
+        raise NotImplementedError(
+            f"The make() method not implemented for entry {self.identifier}"
+        )

ml-agents-envs/mlagents_envs/registry/binary_utils.py

@@ -0,0 +1,193 @@
+import urllib.request
+import tempfile
+import os
+import uuid
+import shutil
+import glob
+import yaml
+
+from zipfile import ZipFile
+from sys import platform
+from typing import Tuple, Optional, Dict, Any
+
+# The default logical block size is 8192 bytes (8 KB) for UFS file systems.
+BLOCK_SIZE = 8192
+
+
+def get_local_binary_path(name: str, url: str) -> str:
+    """
+    Returns the path to the executable previously downloaded with the name argument. If
+    None is found, the executable at the url argument will be downloaded and stored
+    under name for future uses.
+    :param name: The name that will be given to the folder containing the extracted data
+    :param url: The URL of the zip file
+    """
+    path = get_local_binary_path_if_exists(name)
+    if path is None:
+        print(f"Local environment {name} not found, downloading environment from {url}")
+        download_and_extract_zip(url, name)
+    path = get_local_binary_path_if_exists(name)
+    if path is None:
+        raise FileNotFoundError("Binary not found")
+    return path
+
+
+def get_local_binary_path_if_exists(name: str) -> Optional[str]:
+    """
+    Recursively searches for a Unity executable in the extracted files folders. This is
+    platform dependent : It will only return a Unity executable compatible with the
+    computer's OS. If no executable is found, None will be returned.
+    """
+    _, bin_dir = get_tmp_dir()
+    extension = None
+
+    if platform == "linux" or platform == "linux2":
+        extension = "*.x86_64"
+    if platform == "darwin":
+        extension = "*.app"
+    if platform == "win32":
+        extension = "*.exe"
+    if extension is None:
+        raise NotImplementedError("No extensions found for this platform.")
+    path = os.path.join(bin_dir, name, "**", extension)
+    candidates = glob.glob(path, recursive=True)
+    if len(candidates) == 0:
+        return None
+    else:
+        return candidates[0]
+
+
+def get_tmp_dir() -> Tuple[str, str]:
+    """
+    Returns the path to the folder containing the downloaded zip files and the extracted
+    binaries. If these folders do not exist, they will be created.
+    :retrun: Tuple containing path to : (zip folder, extracted files folder)
+    """
+    MLAGENTS = "ml-agents-binaries"
+    ZIP_FOLDER_NAME = "tmp_zip"
+    BINARY_FOLDER_NAME = "binaries"
+    mla_directory = os.path.join(tempfile.gettempdir(), MLAGENTS)
+    if not os.path.exists(mla_directory):
+        os.makedirs(mla_directory)
+    zip_directory = os.path.join(tempfile.gettempdir(), MLAGENTS, ZIP_FOLDER_NAME)
+    if not os.path.exists(zip_directory):
+        os.makedirs(zip_directory)
+    bin_directory = os.path.join(tempfile.gettempdir(), MLAGENTS, BINARY_FOLDER_NAME)
+    if not os.path.exists(bin_directory):
+        os.makedirs(bin_directory)
+    return (zip_directory, bin_directory)
+
+
+def download_and_extract_zip(url: str, name: str) -> None:
+    """
+    Downloads a zip file under a URL, extracts its contents into a folder with the name
+    argument and gives chmod 755 to all the files it contains. Files are downloaded and
+    extracted into special folders in the temp folder of the machine.
+    :param url: The URL of the zip file
+    :param name: The name that will be given to the folder containing the extracted data
+    """
+    zip_dir, bin_dir = get_tmp_dir()
+    binary_path = os.path.join(bin_dir, name)
+    if os.path.exists(binary_path):
+        shutil.rmtree(binary_path)
+
+    # Download zip
+    try:
+        request = urllib.request.urlopen(url)
+    except urllib.error.HTTPError as e:  # type: ignore
+        e.msg += " " + url
+        raise
+    zip_size = int(request.headers["content-length"])
+    zip_file_path = os.path.join(zip_dir, str(uuid.uuid4()) + ".zip")
+    with open(zip_file_path, "wb") as zip_file:
+        downloaded = 0
+        while True:
+            buffer = request.read(BLOCK_SIZE)
+            if not buffer:
+                # There is nothing more to read
+                break
+            downloaded += len(buffer)
+            zip_file.write(buffer)
+            downloaded_percent = downloaded / zip_size * 100
+            print_progress(f"  Downloading {name}", downloaded_percent)
+        print("")
+
+    # Extraction
+    with ZipFileWithProgress(zip_file_path, "r") as zip_ref:
+        zip_ref.extract_zip(f"  Extracting  {name}", binary_path)  # type: ignore
+    print("")
+
+    # Clean up zip
+    print_progress(f"  Cleaning up {name}", 0)
+    os.remove(zip_file_path)
+
+    # Give permission
+    for f in glob.glob(binary_path + "/**/*", recursive=True):
+        # 16877 is octal 40755, which denotes a directory with permissions 755
+        os.chmod(f, 16877)
+    print_progress(f"  Cleaning up {name}", 100)
+    print("")
+
+
+def print_progress(prefix: str, percent: float) -> None:
+    """
+    Displays a single progress bar in the terminal with value percent.
+    :param prefix: The string that will precede the progress bar.
+    :param percent: The percent progression of the bar (min is 0, max is 100)
+    """
+    BAR_LEN = 20
+    percent = min(100, max(0, percent))
+    bar_progress = min(int(percent / 100 * BAR_LEN), BAR_LEN)
+    bar = "|" + "\u2588" * bar_progress + " " * (BAR_LEN - bar_progress) + "|"
+    str_percent = "%3.0f%%" % percent
+    print(f"{prefix} : {bar} {str_percent} \r", end="", flush=True)
+
+
+def load_remote_manifest(url: str) -> Dict[str, Any]:
+    """
+    Converts a remote yaml file into a Python dictionary
+    """
+    tmp_dir, _ = get_tmp_dir()
+    try:
+        request = urllib.request.urlopen(url)
+    except urllib.error.HTTPError as e:  # type: ignore
+        e.msg += " " + url
+        raise
+    manifest_path = os.path.join(tmp_dir, str(uuid.uuid4()) + ".yaml")
+    with open(manifest_path, "wb") as manifest:
+        while True:
+            buffer = request.read(BLOCK_SIZE)
+            if not buffer:
+                # There is nothing more to read
+                break
+            manifest.write(buffer)
+    try:
+        result = load_local_manifest(manifest_path)
+    finally:
+        os.remove(manifest_path)
+    return result
+
+
+def load_local_manifest(path: str) -> Dict[str, Any]:
+    """
+    Converts a local yaml file into a Python dictionary
+    """
+    with open(path) as data_file:
+        return yaml.safe_load(data_file)
+
+
+class ZipFileWithProgress(ZipFile):
+    """
+    This is a helper class inheriting from ZipFile that allows to display a progress
+    bar while the files are being extracted.
+    """
+
+    def extract_zip(self, prefix: str, path: str) -> None:
+        members = self.namelist()
+        path = os.fspath(path)
+        total = len(members)
+        n = 0
+        for zipinfo in members:
+            self.extract(zipinfo, path, None)  # type: ignore
+            n += 1
+            print_progress(prefix, n / total * 100)

ml-agents-envs/mlagents_envs/registry/remote_registry_entry.py

@@ -0,0 +1,62 @@
+from sys import platform
+from typing import Optional, Any
+from mlagents_envs.environment import UnityEnvironment
+from mlagents_envs.base_env import BaseEnv
+from mlagents_envs.registry.binary_utils import get_local_binary_path
+from mlagents_envs.registry.base_registry_entry import BaseRegistryEntry
+
+
+class RemoteRegistryEntry(BaseRegistryEntry):
+    def __init__(
+        self,
+        identifier: str,
+        expected_reward: float,
+        description: str,
+        linux_url: Optional[str],
+        darwin_url: Optional[str],
+        win_url: Optional[str],
+    ):
+        """
+        A RemoteRegistryEntry is an implementation of BaseRegistryEntry that uses a
+        Unity executable downloaded from the internet to launch a UnityEnvironment.
+        __Note__: The url provided must be a link to a `.zip` file containing a single
+        compressed folder with the executable inside. There can only be one executable
+        in the folder and it must be at the root of the folder.
+        :param identifier: The name of the Unity Environment.
+        :param expected_reward: The cumulative reward that an Agent must receive
+        for the task to be considered solved.
+        :param description: A description of the Unity Environment. Contains human
+        readable information about potential special arguments that the make method can
+        take as well as information regarding the observation, reward, actions,
+        behaviors and number of agents in the Environment.
+        :param linux_url: The url of the Unity executable for the Linux platform
+        :param darwin_url: The url of the Unity executable for the OSX platform
+        :param win_url: The url of the Unity executable for the Windows platform
+        """
+        super().__init__(identifier, expected_reward, description)
+        self._linux_url = linux_url
+        self._darwin_url = darwin_url
+        self._win_url = win_url
+
+    def make(self, **kwargs: Any) -> BaseEnv:
+        """
+        Returns the UnityEnvironment that corresponds to the Unity executable found at
+        the provided url. The arguments passed to this method will be passed to the
+        constructor of the UnityEnvironment (except for the file_name argument)
+        """
+        url = None
+        if platform == "linux" or platform == "linux2":
+            url = self._linux_url
+        if platform == "darwin":
+            url = self._darwin_url
+        if platform == "win32":
+            url = self._win_url
+        if url is None:
+            raise FileNotFoundError(
+                f"The entry {self.identifier} does not contain a valid url for this "
+                "platform"
+            )
+        path = get_local_binary_path(self.identifier, url)
+        if "file_name" in kwargs:
+            kwargs.pop("file_name")
+        return UnityEnvironment(file_name=path, **kwargs)