Add sample showing how to integrate an Anylogic sim

This sample leverages Baobab.

.gitignore

@@ -151,3 +151,6 @@ etc/
 
 # Ray checkpoints
 **/checkpoint_*/
+
+# Exported anylogic sim
+exported/

.pre-commit-config.yaml

@@ -23,3 +23,4 @@ repos:
     rev: v1.1.1
     hooks:
     -   id: mypy
+        additional_dependencies: ['types-requests==2.31.0.1']

examples/getting-started-anylogic-sim/README.md

@@ -0,0 +1,112 @@
+# Run an Anylogic sim in AML
+
+In this folder we show how to get started training an RL agent on Azure ML
+with an Anylogic sim.
+
+### What this sample covers
+
+- Integration of a custom Gymnasium simulation environment with RLlib
+- How to train an agent using the simulation environment on AML
+
+### What this sample does not cover
+
+- How to optimize the training algorithm for best performance
+- How to visualize the performance of the agent
+- How to evaluate the agent
+- How to deploy the agent
+
+## Prerequisites
+
+- Install the Azure CLI on your machine:
+```
+pip install azure-cli
+```
+- Add the ML extension:
+```
+az extension add -n ml
+```
+- [Create an AML workspace and compute cluster](https://azure.github.io/plato/#create-azure-resources)
+- Create an AML environment using the docker instructions provided in the
+  ``docker`` folder. Go in that folder and run the following command:
+```bash
+az ml environment create --name anylogic-env --build-context ./ --dockerfile-path Dockerfile --resource-group $YOUR_RESOURCE_GROUP --workspace-name $YOUR_WORKSPACE
+```
+- Have an Anylogic sim that supports Bonsai integration. See
+  [this](https://github.com/microsoft/bonsai-anylogic) for details.
+
+## Prepare the Anylogic sim
+
+This sample was built around the sim that you can find
+[here](https://github.com/microsoft/bonsai-anylogic/tree/master/samples/abca).
+Unzip the content of ``export.zip`` to this location: ``src/export``. Make
+sure that you have the script for launching the sim in a Linux environment
+and that the name of the script ends with ``_linux.sh``.
+
+## Test Locally
+
+If you have Docker available and working on your machine, you could try to
+test if everything starts up correctly on your machine before sending the job
+to AML.
+The following instructions are made for a Unix-like system. You'll need to
+modify them to run on Windows.
+To do this:
+
+1. Go into the ``docker`` folder and build the docker image:
+
+```bash
+docker build -t anylogic-sim .
+```
+
+2. Go to the sample folder and run the sample:
+
+```bash
+docker run --rm -it -e LOG_LEVEL=debug -v $(pwd)/src:/opt/src anylogic-sim bash /opt/src/start.sh
+```
+
+After some time you should see that training starts as the terminal will show
+the communication between the sim, Baobab and RLlib. When that happens, you
+can proceed to the next section. If you see errors, please try to fix them
+before proceeding further.
+
+## Tutorial: Run on AML
+
+After you checked that the simulation works properly, follow these steps to
+train an RL agent on AML:
+
+1. Modify the ``job.yml`` file by changing the name of the AML ``environment``
+   and ``compute`` to be the same as those you created in the prerequisites
+   section.
+
+2. Launch the job using the Azure CLI:
+```
+az ml job create -f job.yml --workspace-name $YOUR_WORKSPACE --resource-group $YOUR_RESOURCE_GROUP
+```
+
+3. Check that it is running by finding the job you launched in [AML
+   studio](https://ml.azure.com/). You should see that ``ray`` is writing
+   logs in the *Outputs + logs* tab in the ``user_logs`` folder.
+
+4. Once the job is completed, the model checkpoints can also be found in AML
+   studio under the *Outputs + Logs* tab of your job
+   in the ``outputs`` folder.
+
+## Modify the sample to use your own sim
+
+If you have your own Anylogic sim, there are a couple of changes you have to
+make.
+
+In ``src/sim.py``:
+
+- In ``SimWrapper.__init__``: modify ``self.action_space`` to be the possible
+  actions, modify ``self.observation_space`` to contain all states that the
+  agent has access to, and modify ``self.config`` in order to pass the config
+  needed by the sim.
+- Modify ``SimWrapper.reward`` in ``src/sim.py`` by adapting it to the
+  problem you want to solve.
+- Modify ``SimWrapper.terminal`` in ``src/sim.py`` to return ``True`` when
+  conditions are met to terminate the episode
+- Modify ``SimWrapper.truncate`` in ``src/sim.py``
+
+To understand the meaning of ``reward``, ``terminal``, and ``truncate`` we
+suggest checking Gymnasium's
+[documentation](https://gymnasium.farama.org/api/env/#gymnasium.Env.step).

examples/getting-started-anylogic-sim/docker/Dockerfile

@@ -0,0 +1,12 @@
+FROM mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu22.04
+
+RUN apt-get update \
+	# Otherwise installing default-jre-headless fails
+	&& mkdir -p /usr/share/man/man1/ \
+	&& DEBIAN_FRONTEND=noninteractive apt-get install -y default-jre-headless git memcached \
+	&& apt clean \
+	&& rm -rf /var/lib/apt/lists/*
+
+WORKDIR /opt/src
+COPY requirements.txt /opt/src
+RUN pip install -Ur requirements.txt

examples/getting-started-anylogic-sim/docker/requirements.txt

@@ -0,0 +1,11 @@
+aiocache[memcached]==0.12.1
+azure-ai-ml==1.7.2
+azureml-core==1.51.0
+azureml-defaults==1.51.0
+azureml-mlflow==1.51.0
+fastapi==0.95.2
+git+https://github.com/Azure/plato.git@main
+gunicorn==20.1.0
+ray[air,default,rllib]==2.4.0
+ray_on_aml==0.2.4
+torch==2.0.1

examples/getting-started-anylogic-sim/job.yml

@@ -0,0 +1,15 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
+code: src
+command: >-
+  ./start.sh
+environment: azureml:anylogic-env@latest
+compute: azureml:env-medium
+display_name: anylogic-sim
+experiment_name: anylogic-sim
+description: Run a Bonsai Anylogic sim on AML.
+# Needed for using ray on AML
+distribution:
+   type: mpi
+# Modify the following and num_rollout_workers in main to use more workers
+resources:
+  instance_count: 1

examples/getting-started-anylogic-sim/src/main.py

@@ -0,0 +1,61 @@
+"""
+Main script for training an agent on Azure ML.
+
+You can test this code locally by launching the script with the flag
+``--test-local``
+"""
+import argparse
+import sys
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.tune.logger import pretty_print
+from ray.tune.registry import register_env
+from ray_on_aml.core import Ray_On_AML
+from sim import SimWrapper as SimEnv
+
+# Register the simulation as an RLlib environment.
+register_env("sim_env", lambda config: SimEnv(config))
+
+
+def train(local=False):
+    # Define the algo for training the agent
+    algo = (
+        PPOConfig()
+        # Modify also instance_count in job definition to use more than one worker
+        # Setting workers to zero allows using breakpoints in sim for debugging
+        .rollouts(num_rollout_workers=1 if not local else 0)
+        .resources(num_gpus=0)
+        # Set the training batch size to the appropriate number of steps
+        .training(train_batch_size=4_000)
+        .environment(env="sim_env")
+        .build()
+    )
+    # Train for 10 iterations
+    for i in range(10):
+        result = algo.train()
+        print(pretty_print(result))
+
+    # outputs can be found in AML Studio under the "Outputs + Logs" tab of your job
+    checkpoint_dir = algo.save(checkpoint_dir="./outputs")
+    print(f"Checkpoint saved in directory {checkpoint_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--test-local", action="store_true", default=False)
+    args = parser.parse_args()
+
+    if args.test_local:
+        train(args.test_local)
+        sys.exit()
+
+    ray_on_aml = Ray_On_AML()
+    ray = ray_on_aml.getRay()
+
+    if ray:
+        print("head node detected")
+        ray.init(address="auto")
+        print(ray.cluster_resources())
+        train(args.test_local)
+    else:
+        print("in worker node")

examples/getting-started-anylogic-sim/src/sim.py

@@ -0,0 +1,63 @@
+import json
+
+import requests
+from gymnasium import Env, spaces
+
+from platotk.logger import log
+from platotk.serialize import GymEncoder, check_and_transform
+
+BASE_URL = "http://localhost:8000"
+
+
+class SimWrapper(Env):
+    def __init__(self, env_config):
+        self.base_url = BASE_URL
+
+        self.action_space = spaces.Dict(
+            {
+                "numResourceA": spaces.Discrete(20, start=1),
+                "numResourceB": spaces.Discrete(20, start=1),
+                "processTime": spaces.Box(1.0, 12.0),
+                "conveyorSpeed": spaces.Box(0.01, 1.0),
+            }
+        )
+        self.observation_space = spaces.Dict({"arrivalRate": spaces.Box(0.5, 2.0)})
+        self.config = {
+            "arrivalRate": 0.5,
+            "sizeBufferQueues": 45,
+        }
+
+    def reset(self, *, seed=None, options=None):
+        log.debug("Reset send.")
+        resp = requests.post(self.base_url + "/reset", json={"config": self.config})
+        state = resp.json()
+        log.debug("Reset response.")
+        return check_and_transform(self.observation_space, state), {}
+
+    def step(self, action):
+        log.debug("Send step.")
+        json_action = json.dumps({"action": action}, cls=GymEncoder)
+        resp = requests.post(
+            self.base_url + "/step",
+            data=json_action,
+            headers={"Content-Type": "application/json"},
+        )
+        state = resp.json()
+        log.debug("Step response.")
+
+        return (
+            check_and_transform(self.observation_space, state),
+            self.reward(state),
+            self.terminal(state),
+            self.truncate(state),
+            {},
+        )
+
+    def reward(self, state):
+        return -state.get("costPerProduct") - 1000 * state.get("exceededCapacityFlag")
+
+    def terminal(self, state):
+        return state.get("exceededCapacityFlag") == 1 or state.get("simTimeMonths") >= 6
+
+    def truncate(self, state):
+        return self.terminal(state)

examples/getting-started-anylogic-sim/src/start.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/sh
+set -xe
+memcached -p 11211 -u memcache -l 127.0.0.1 -d
+
+gunicorn \
+	--worker-class uvicorn.workers.UvicornWorker \
+	--bind '127.0.0.1:8000' \
+	platotk.baobab:app &
+
+export SIM_API_HOST=http://localhost:8000
+export SIM_CONTEXT={}
+export SIM_WORKSPACE=dummy
+export SIM_ACCESS_KEY=dummy
+
+bash "$(find -name '*_linux.sh')" &
+python -u main.py --test-local

requirements/dev_requirements.in

@@ -3,3 +3,4 @@ flake8
 isort
 mypy
 pytest
+types-requests

requirements/dev_requirements.txt

@@ -47,5 +47,9 @@ tomli==2.0.1
     #   black
     #   mypy
     #   pytest
+types-requests==2.31.0.1
+    # via -r requirements/dev_requirements.in
+types-urllib3==1.26.25.13
+    # via types-requests
 typing-extensions==4.5.0
     # via mypy