Jane Street1h 13m
Playing Atari Games in OCaml
Conference Presentation, Tutorial
The presentation demonstrates applying machine learning (ML) to Atari games using OCaml, leveraging functional programming and type safety while interfacing with Python-based ML frameworks (PyTorch/TensorFlow) for GPU acceleration.
Supervised Learning Segment:
- In supervised learning, the goal is to minimize a loss function by optimizing parameters $\theta$ (continuous space) rather than searching a discrete set of algorithms.
- Gradient descent is employed: parameters are initialized randomly, predictions are made, loss is computed against ground truth labels, and parameters are updated via the gradient scaled by a learning rate ($\alpha$).
- Data is split into training and validation sets to prevent overfitting on the training distribution.
- Deep learning models achieve >99.5% accuracy on image classification (e.g., distinguishing pythons from camels), surpassing human performance.
- Architecture Evolution:
- Simple linear models (2 layers, ~30M parameters) struggle with complex images; Convolutional Neural Networks (CNNs) are required to exploit spatial locality (edge detection $\to$ higher-level features like wheels $\to$ objects).
- Non-linear activation functions (e.g., ReLU) are used instead of sigmoids for computational efficiency and to prevent the "vanishing gradient" problem.
- ImageNet Progress: Error rates dropped from 26% (2001) to 3.1% (2016), with deep learning models (AlexNet, 2012) surpassing human capability by 2015.
- OCaml Implementation: OCaml types are used to define the model structure (input dimensions, layer sizes), while PyTorch handles tensor operations and automatic differentiation (backpropagation) via a C++ backend called from OCaml.
- Training Loop: Uses mini-batches (e.g., size 64) to fit data into GPU memory (e.g., 8GB RAM on RTX 2070); models require hours to days of GPU training.
Reinforcement Learning (RL) Segment:
- RL involves an agent interacting with an environment, receiving states ($s_t$), taking actions ($a_t$), and receiving rewards ($r_t$) to maximize the cumulative discounted reward ($\gamma \approx 0.99$).
- Deep Q-Learning (DQN):
- Uses the Bellman equation to define the loss: $L = \mathbb{E}[(Q(s,a) - (r + \gamma \max_{a'} Q(s',a')))^2]$.
- The Q-network approximates the optimal Q-function ($Q^*$), taking image frames (often 4 stacked frames to capture motion) as input and outputting Q-values for all possible actions.
- An $\epsilon$-greedy policy balances exploration (random actions) and exploitation (taking the best predicted action), with $\epsilon$ decreasing over time.
- Game-Specific Observations:
- Pong: A simple 2-layer linear model suffices; the agent learns to return the ball but struggles with "inaction" (requires rapid up/down toggling).
- Breakout: A CNN is required; the agent discovers a "tunneling" strategy (digging holes in brick layers) to clear the board more efficiently.
- Space Invaders: The agent learns to shoot invaders but initially fails to prioritize specific high-value targets or self-preservation until extensive training occurs.
- Limitations: Training requires hundreds to thousands of episodes (days on GPU) compared to human learning (minutes), due to the lack of prior world physics knowledge.
- Alternative Algorithms: Policy gradient methods (e.g., A3C, PPO) are noted as trending alternatives that learn faster by running multiple environments in parallel on CPU cores.
Technical Implementation Details:
- Interoperability: OCaml uses
pymlto interface with Python libraries (OpenAI Gym, PyTorch) without duplicating memory, allowing data to be shared directly between OCaml and Python. - Ecosystem Constraints: The OCaml ML ecosystem is smaller than Python's; the speaker recommends mimicking Python API naming conventions to allow users to find existing Python documentation.
- Hardware: Experiments were conducted on consumer-grade hardware (single RTX 2070, 8GB VRAM); no multi-GPU or distributed training was utilized.
- Fairness in AI vs. Human: In RTS games (e.g., StarCraft), AI is restricted by Actions Per Minute (APM) limits to match human capabilities, otherwise, AI exploits "micro-management" precision bugs impossible for humans to counter.
- Interoperability: OCaml uses
Future Directions & Challenges:
- Transfer Learning & Self-Supervision: Moving toward models that learn general world physics (self-supervised) to reduce the need for massive labeled datasets.
- Generalization: Research is ongoing to enable agents trained on one Atari game (e.g., Breakout) to adapt to variations (e.g., different paddle shapes) without retraining from scratch.
- Market Applications: RL is applicable to markets with noisy time-series data and the need for perfect simulators, though current challenges include data noise and the difficulty of modeling real-world physics accurately.
- Open Source: The demonstrated PyTorch and TensorFlow OCaml bindings are available on GitHub for community feedback and extension.