Introduction

Recently I started participating in Reinforcement Learning (RL) competitions hosted by ArenaX Labs on the SAI Platform. SAI is a platform that that hosts competitions and provides Gymnasium based environments for the RL community to compete and advance their research.

The competition I've been participating on is based on the Franka Golf Environment. The purpose of the competition is to train an RL agent to pick up a golf club and striking and successfully putting the ball in the designated hole.

Figure 1: Franka Golf Environment provided by the SAI platform.

Reproducibility and Docker

RL research often involves a complex stack of software ranging from different frameworks (Pytorch, Tensorflow, etc), the environment, various dependencies and the RL agent code to name a few. As such, reproducible docker environments become very important and Docker is a great tool for that. This post provides a simple docker template to quickly get setup for the RL experimentation.

The template can be found at this github repo. This repo contains three essential files that define all the dependencies required to set up the environment. These files are named Dockerfile, docker-compose.yml and devcontainer.json.

Dockerfile

This file essentially defines the image that the docker containers will be running. It extends from the stablebaselines3 image and installs the sai-rl and sai-mujoco package. The snippet below shows the contents of the dockerfile.

FROM stablebaselines/rl-baselines3-zoo:latest

# Install sai-rl
# Original container uses micromamba and uv
RUN \
    eval "$(micromamba shell hook --shell bash)" && \
    micromamba activate && \
    uv pip install --system sai-rl && \
    uv pip install --system sai-mujoco && \
    uv cache clean

WORKDIR /workspace

CMD ["/bin/bash"]

docker-compose.yml

This file defines the service to build and run the docker image in a container. This files also contains other instructions related to the container such as volume information, network config, runtime config, etc. The snippet below shows the contents of the docker compose file.

version: "3.8"

services:
  franka-golf-saad:
    build: .
    image: saad1998/franka-golf-saad
    container_name: franka-golf-saad
    working_dir: /workspace
    user: root
    volumes:
      - ./:/workspace
    tty: true
    stdin_open: true
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    network_mode: host
    ipc: host
    command: bash
    restart: "no"

devcontainer.json

This file contains the configurations needed by the devcontainers extension in VSCode. When using VSCode, it should automatically detect the .devcontainer configuration if it's present. To reopen the folder in the devcontainer, follow these steps:

1. In the VSCode window, press Ctrl + Shift + P to open the Command Palette.
2. Search for and select Dev Containers: Reopen in Container.

VSCode will:

• Build the Docker container (if it's not built already).
• Open the project in a new window with the environment defined in the container.

Note: The first time it may take a few minutes to build the Docker image and set up the environment.

The snippet below shows the contents of the devcontainer.json file.

{
  "name": "franka-golf-saad",
  "dockerComposeFile": "../docker-compose.yml",
  "service": "franka-golf-saad",
  "workspaceFolder": "/workspace",
  "remoteUser": "root",
  "features": {
    "ghcr.io/devcontainers-extra/features/black:2": {},
    "ghcr.io/devcontainers/features/git:1": {}
  },
  "customizations": {
    "vscode": {
      "settings": {
        "python.defaultInterpreterPath": "/opt/conda/bin/python"
      },
      "extensions": [
        "ms-python.black-formatter"
      ]
    }
  }
}

Conclusion

Creating a reproducible environment is a critical step in effective RL research and development. By using a Docker template, researchers and developers can eliminate the common headaches associated with dependency management and version conflicts. The combination of a Dockerfile, docker-compose.yml, and devcontainer.json provides a robust and portable solution. This setup not only ensures experiments run consistently across different machines but also simplifies collaboration, allowing teams to focus on the core task: training better RL agents.

Whether a person is tackling the Franka Golf Environment or a different challenge, this template offers a solid foundation. It can be easily adapted to specific needs by modifying the dependencies and configurations. This approach helps to accelerate the workflow, making it easier to share results and contribute to the broader RL community.