rheehot / haiku-scalable-example

Scalable distributed reinforcement learning agents on kubernetes

haiku-scalable-example

Scalable reinforcement learning agents on container orchestration

1. Purpose of the project

Implement scalable reinforcement learning agent on the container orchestraion system like k8s.

2. Container Orchestraion

• Kubernetes
• Google Cloud Platform

3. Reinforcement Learning Algorithms

• IMPALA
• A3C
• TBD

4. Architecture

This example will introduce a clear way to deploy scalable reinforcement learning agents to the computing clusters.

5. Install

$ git clone https://github.com/chris-chris/haiku-scalable-example
$ cd haiku-scalable-example
$ pip install -r requirements.txt

6. Execute

v1. Learner + Multi Actor IMPALA wiring through gRPC.

$ python learner_server.py

$ GRPC_HOST=localhost:50051 python actor_client.py &
$ GRPC_HOST=localhost:50051 python actor_client.py &

v2. 1 Learner + Multi Actor IMPALA wiring through gRPC on docker VMs.

prepare

$ docker pull chrisai/haiku-scalable-example-learner:test
$ docker pull chrisai/haiku-scalable-example-actor:test

$ docker network create --subnet 172.20.0.0/16 --ip-range 172.20.240.0/20 multi-host-network

run

$ docker run -d -p 127.0.0.1:50051:50051 --network=multi-host-network --ip=172.20.240.1 chrisai/haiku-scalable-example-learner:test
$ docker run -d --env GRPC_HOST=172.20.240.1:50051 --network=multi-host-network chrisai/haiku-scalable-example-actor:test

wanna see logs?

$ docker ps
$ docker attach [CONTAINER ID]

v3. 1 Learner + Multi Actor IMPALA wiring through gRPC on k8s.

• Install minikube

https://kubernetes.io/docs/tasks/tools/install-minikube/

• Run

$ kubectl apply -f impala.yml

• Wanna see logs?

$ kubectl logs -f impala learner
$ kubectl logs -f impala actor

7. To-dos

• v1. 1 Learner + Multi Actor IMPALA wiring through gRPC.
• v2. 1 Learner + Multi Actor IMPALA wiring through gRPC on docker VMs.
• v3. 1 Learner + Multi Actor IMPALA wiring through gRPC on k8s.
• Optimize the model weight serialization for the performance.
• v4. Multi Learner + Multi Actor IMPALA wiring through gRPC on k8s.
• Implement other distributed RL algorithms
• Asynchronous Processing via Queue
• Monitor the computing resource usages

8. Reference

I used Deepmind's open sources haiku, rlax, and google jax

• https://github.com/google/jax
• https://github.com/deepmind/rlax
• https://github.com/deepmind/haiku
• https://github.com/kubernetes/kubernetes
• https://github.com/kent-williams/grpc-python-kubernetes