[WIP] Unity Environment Registry

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

@@ -0,0 +1,4 @@
+from mlagents_envs.registry.unity_env_registry import UnityEnvRegistry  # noqa F401
+
+# Important : This loads all the registry entries to default environments when importing
+import mlagents_envs.registry.default_envs  # noqa F401

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,154 @@
+import urllib.request
+import tempfile
+import os
+import uuid
+import shutil
+import glob
+
+from zipfile import ZipFile
+from sys import platform
+from typing import Tuple, Optional
+
+
+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
+    """
+    # The default logical block size is 8192 bytes (8 KB) for UFS file systems.
+    BLOCK_SIZE = 8192
+    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
+    request = urllib.request.urlopen(url)
+    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)
+
+
+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/default_envs.py

@@ -0,0 +1,78 @@
+from mlagents_envs.registry.unity_env_registry import UnityEnvRegistry
+from mlagents_envs.registry.remote_registry_entry import RemoteRegistryEntry
+
+
+registry = UnityEnvRegistry()
+
+# 3dBall
+registry.register(
+    RemoteRegistryEntry(
+        identifier="3DBall",
+        expected_reward=100,
+        description="""
+## 3DBall: 3D Balance Ball
+- Set-up: A balance-ball task, where the agent balances the ball on it's head.
+- Goal: The agent must balance the ball on it's head for as long as possible.
+- Agents: The environment contains 12 agents of the same kind, all using the
+  same Behavior Parameters.
+- Agent Reward Function:
+  - +0.1 for every step the ball remains on it's head.
+  - -1.0 if the ball falls off.
+- Behavior Parameters:
+  - Vector Observation space: 8 variables corresponding to rotation of the agent
+    cube, and position and velocity of ball.
+  - Vector Observation space (Hard Version): 5 variables corresponding to
+    rotation of the agent cube and position of ball.
+  - Vector Action space: (Continuous) Size of 2, with one value corresponding to
+    X-rotation, and the other to Z-rotation.
+  - Visual Observations: None.
+- Float Properties: Three
+  - scale: Specifies the scale of the ball in the 3 dimensions (equal across the
+    three dimensions)
+    - Default: 1
+    - Recommended Minimum: 0.2
+    - Recommended Maximum: 5
+  - gravity: Magnitude of gravity
+    - Default: 9.81
+    - Recommended Minimum: 4
+    - Recommended Maximum: 105
+  - mass: Specifies mass of the ball
+    - Default: 1
+    - Recommended Minimum: 0.1
+    - Recommended Maximum: 20
+- Benchmark Mean Reward: 100
+        """,
+        linux_url=None,
+        darwin_url="https://storage.googleapis.com/mlagents-test-environments/1.0.0/darwin/3DBall.zip",  # noqa: E501
+        win_url=None,
+    )
+)
+
+# Basic
+registry.register(
+    RemoteRegistryEntry(
+        identifier="Basic",
+        expected_reward=0.93,
+        description="""
+## Basic
+- Set-up: A linear movement task where the agent must move left or right to
+  rewarding states.
+- Goal: Move to the most reward state.
+- Agents: The environment contains one agent.
+- Agent Reward Function:
+  - -0.01 at each step
+  - +0.1 for arriving at suboptimal state.
+  - +1.0 for arriving at optimal state.
+- Behavior Parameters:
+  - Vector Observation space: One variable corresponding to current state.
+  - Vector Action space: (Discrete) Two possible actions (Move left, move
+    right).
+  - Visual Observations: None
+- Float Properties: None
+- Benchmark Mean Reward: 0.93
+        """,
+        linux_url=None,
+        darwin_url="https://storage.googleapis.com/mlagents-test-environments/1.0.0/darwin/Basic.zip",  # noqa: E501
+        win_url=None,
+    )
+)

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)