mscrnt/diambra-trainer repository overview

⁠Diambra Docker (Stable Baselines 3)

This Docker image provides a training environment for DIAMBRA Arena⁠ with Stable Baselines 3⁠ (SB3) and PyTorch support. It includes TensorBoard for real-time monitoring of your reinforcement learning training and uses Docker-in-Docker (DinD) to support multiple environments.

Source Code (GitHub)⁠

⁠Table of Contents

• Diambra Docker Image⁠
  • Table of Contents⁠
  • Features⁠
  • Prerequisites and Preparing Directories⁠
    • Directory Setup⁠
    • Option 1: Use Existing DIAMBRA Credentials⁠
    • Option 2: Generate a DIAMBRA Credentials File via Docker⁠
  • ⚠️ Warning⁠
  • Usage⁠
    • Pulling the Docker Image⁠
    • Running the Container⁠
  • Environment Variables⁠
  • Exposed Ports⁠
  • Accessing TensorBoard⁠
  • License⁠
  • Contributing⁠

⁠Features

• Stable Baselines 3 Integration: Fully integrated with DIAMBRA Arena for reinforcement learning tasks.
• PyTorch 2.4.1 with CUDA 12.4: GPU acceleration for efficient training.
• DIAMBRA CLI: Simplified environment setup and management.
• TensorBoard: Real-time monitoring of training metrics.
• Docker-in-Docker (DinD): Allows for running Docker containers within the Docker environment.

⁠Prerequisites and Preparing Directories

Before running the container, you need to set up your directories and provide DIAMBRA credentials. You can either use an existing credentials file or generate a new one inside the Docker container.

⁠Directory Setup

You will need the following directories on your host machine, which will be mounted into the Docker container:

1. ROMs Directory: Contains your game ROMs.

   mkdir -p /path/to/roms
   

   Copy

2. Scripts Directory: Contains your training scripts (e.g., train.py).

   mkdir -p /path/to/scripts
   

   Copy

3. DIAMBRA Credentials: You can either copy an existing credentials file or generate one inside the container.

   mkdir -p /path/to/diambra
   

   Copy

4. Output Directory: Stores trained models and logs.

   mkdir -p /path/to/output
   

   Copy

⁠Option 1: Use Existing DIAMBRA Credentials

If you already have a DIAMBRA credentials file, follow these steps:

1. Copy your credentials file into the diambra directory on your host machine:

   cp /your/credentials/file /path/to/diambra/credentials
   

   Copy

2. When running the container, mount the credentials file by adding the -v /path/to/diambra:/workspace/diambra option to your Docker command.

⁠Option 2: Generate a DIAMBRA Credentials File via Docker

If you don’t have a credentials file, the container will enter Setup Mode when it doesn’t detect the file. Follow these steps inside the container to generate the credentials:

1. Start the container (If you haven't already):

   docker run --privileged -v /path/to/diambra:/workspace/diambra -it mscrnt/diambra-trainer:latest
   

   Copy

2. Inside the container, start the Docker daemon:

   sudo dockerd > /dev/null 2>&1 &
   

   Copy

3. Generate the credentials file using the DIAMBRA CLI:

   diambra run -n --path.credentials "/workspace/diambra/credentials"
   

   Copy

4. Once the credentials file is generated, you can either run the training script manually or restart the container with the saved credentials.

⁠⚠️ Warning

⁠Usage

⁠Pulling the Docker Image

You can pull the pre-built image from Docker Hub:

docker pull mscrnt/diambra-trainer:latest

⁠Running the Container

Use the following command to run the Docker container:

docker run --privileged \
  -v /path/to/roms:/workspace/roms \
  -v /path/to/scripts:/workspace/scripts \
  -v /path/to/diambra:/workspace/diambra \
  -v /path/to/output:/workspace/output \
  -e SCALE=2 \
  -e EXTRA_ARGS="--batch-size 64" \
  -e TRAINING_SCRIPT="/path/to/your_training_script.py" \
  -p 7007:6006 \
  --name diambra-trainer \
  -it mscrnt/diambra-trainer:latest

• --privileged: Required for DinD to function properly.
• Volume Mounts (-v): Mounts your host directories into the container.
• Environment Variables (-e): Sets the number of DIAMBRA environments, additional arguments, and specifies the training script.
• Port Mapping (-p): Exposes TensorBoard on port 6006 and maps it to 7007 on the host.

⁠Environment Variables

• SCALE: Adjust to change the number of parallel environments (default: 1).

  -e SCALE=4
  

  Copy
• EXTRA_ARGS: Pass additional arguments to your training script (default: "").

  -e EXTRA_ARGS="--learning-rate 0.0001 --batch-size 64"
  

  Copy
• TRAINING_SCRIPT: Specify a custom path for the training script (default: /workspace/scripts/train.py).

  -e TRAINING_SCRIPT="/path/to/your_training_script.py"
  

  Copy
• AUTO: Automatically run the training script on startup (default: true).

  -e AUTO="true"
  

  Copy
• STOP_AFTER_RUN: Stop the container after the training script completes (default: false).

  -e STOP_AFTER_RUN="true"
  

  Copy
• DIAMBRAROMSPATH: Path to the directory containing game ROMs (default: /workspace/roms).
  Note: It is recommended not to change this unless you are familiar with how DIAMBRA Arena manages ROMs.

⁠Exposed Ports

• 6006: TensorBoard web interface.

⁠Accessing TensorBoard

After starting the container, you can access TensorBoard by navigating to:

http://localhost:7007

You can monitor your training progress with real-time visualizations of metrics like loss and reward.

⁠License

This project is licensed under the MIT License⁠.

⁠Contributing

Contributions are welcome! Please open an issue or submit a pull request for any improvements or suggestions.