How realwenlongwang/drama works
1. Purpose and top-level shape
- This repo implements Drama, a model-based reinforcement learning system built around Mamba/Mamba2 sequence models.
- The architecture is layered: experiment entrypoints on top, RL agents and environment wrappers in the middle, model components below, and custom CUDA/Triton sequence operators at the bottom.
2. Training and evaluation entrypoints
- train.py is the main orchestration script for running experiments.
- eval.py handles evaluation-only runs.
- config_files/configure.yaml and the training config files drive runtime behavior, letting the same codebase switch among model backbones and behavior-learning variants.
3. RL control plane
- agents.py likely contains the high-level agent logic that ties together world modeling, action selection, and optimization.
- replay_buffer.py is the storage boundary for collected experience, feeding sequences or transitions back into training.
- tools.py and utils.py provide shared experiment utilities such as configuration handling, logging, scheduling, or metric helpers.
4. Environment subsystem
- envs/ contains task-specific environment adapters:
- my_atari.py
- my_dmc.py
- my_memory_maze.py
- envs/parallel.py suggests vectorized or subprocess-based rollout execution.
- env_wrapper.py likely normalizes observation/action APIs across these environments so the agent can train against multiple benchmark families with the same interface.
5. Model-based RL core
- sub_models/world_models.py is the likely center of the learned latent dynamics / world model.
- sub_models/transformer_model.py indicates an alternative sequence backbone.
- sub_models/attention_blocks.py, sub_models/functions_losses.py, and sub_models/laprop.py support model internals, training losses, and optimization behavior.
- This split suggests a design where the world model can be paired with different sequence architectures and loss heads.
6. Behavior learning / policy side
- The README states the repo supports Actor-Critic and PPO behavior models.
- Those policy/value components likely live in sub_models/ and are coordinated by the agent/training layer rather than being separate standalone services.
- This means the system has two main learning loops: world-model learning and behavior optimization.
7. Mamba/Mamba2 implementation layer
- The mamba_ssm/ package is a vendored or embedded sequence-model library.
- Key areas include:
- modules/ for Mamba blocks, MHA, MLP, and composite sequence blocks
- models/ for higher-level model assembly and configs
- ops/ for selective scan and Triton-based kernels
- distributed/ for tensor-parallel or multi-GPU support
- This package is the main architectural dependency for the “Mamba-enabled” part of Drama.
8. Custom sequence operators and kernel boundary
- csrc/selective_scan/ contains compiled CUDA/C++ extensions for the selective scan operator used by Mamba-style recurrence.
- The presence of forward/backward kernels for fp16, bf16, and fp32 indicates performance-critical training/inference paths are implemented below Python.
- mamba_ssm/ops/selective_scan_interface.py is the Python boundary that calls into these native kernels.
9. Distributed / performance tooling
- mamba_ssm/distributed/ suggests optional tensor parallelism or distributed execution for larger models.
- The repo also includes Triton kernels under mamba_ssm/ops/triton/, which likely cover layer norm, state update, chunk scan, and fused sequence operations.
- Together, these pieces show the system is designed for GPU-heavy training with custom acceleration rather than pure PyTorch reference code.
10. Experiment packaging and runtime
- Dockerfile supports containerized execution, mainly for reproducibility and GPU setup.
- requirements.txt plus the README’s installation notes show the repo depends on a specific PyTorch and mamba-ssm stack.
- config_files/ and the README’s CLI examples imply experiments are configured primarily through YAML plus command-line overrides.
11. Data flow through the system
- Environments produce trajectories or observations.
- The agent collects them into the replay buffer.
- Training samples sequences from the buffer to update the world model and behavior policy.
- The world model uses Mamba/Transformer backbones to predict latent dynamics and/or future observations.
- The policy/value side uses those learned representations to improve action selection.
12. What to show in an architecture diagram
- Draw four main layers: entrypoints/config, RL orchestration, environment + replay, and model backends.
- Under model backends, separate Drama-specific sub_models from the reusable mamba_ssm library and then from native CUDA/Triton ops.
- Highlight the two external interfaces: Gym-like environments on one side and GPU compiled kernels on the other.
Main entry points: agents.py agents.py · parallel.py parallel rollout · world_models.py world models.py · mixer_seq_simple.py mamba models
Components 22
Entrypoints
Experiment launch and configuration
- train.py · Main experiment orchestration script
train.py - eval.py · Evaluation-only execution path
eval.py - configure.yaml · Primary YAML configuration for experiments
config_files/configure.yaml
RL Control
Agent orchestration and shared runtime utilities
- agents.py · High-level agent logic tying together rollout, learning, and optimization
agents.py - replay buffer · Experience storage and sampling boundary
replay_buffer.py - tools.py · Shared experiment helpers
tools.py - utils.py · Common utility functions used across training and evaluation
utils.py
Environments
Task adapters and rollout execution
- env wrapper · Normalizes environment interfaces across tasks
env_wrapper.py - parallel rollout · Parallel or subprocess-based environment execution
envs/parallel.py - Atari env · Atari benchmark environment adapter
envs/my_atari.py - DMC env · DeepMind Control environment adapter
envs/my_dmc.py - Memory maze · Memory maze environment adapter
envs/my_memory_maze.py
Drama Models
Task-specific learning modules
- world models.py · Latent dynamics and world model learning
sub_models/world_models.py - transformer model · Alternative sequence-model backbone
sub_models/transformer_model.py - attention blocks · Attention and sequence building blocks
sub_models/attention_blocks.py - losses · Model loss functions and optimization objectives
sub_models/functions_losses.py - laprop · Optimization support utilities
sub_models/laprop.py
Mamba Stack
Reusable sequence model implementation
- mamba models · Higher-level Mamba model configuration and assembly
mamba_ssm/models/mixer_seq_simple.py - mamba blocks · Core Mamba, Mamba2, attention, and MLP modules
mamba_ssm/modules/block.py - selective scan · Python boundary for selective-scan operators
mamba_ssm/ops/selective_scan_interface.py
Native Kernels
Compiled CUDA/C++ and Triton operators
- triton ops · Triton fused kernels for normalization and state updates
mamba_ssm/ops/triton/ssd_combined.py - selective scan cuda · Compiled CUDA/C++ selective-scan implementation
csrc/selective_scan/selective_scan.cpp
Connections 23
- configure.yaml → train.py: configures · Runtime behavior and model variant selection
- configure.yaml → eval.py: configures · Evaluation settings and overrides
- train.py → agents.py: starts · Launches training orchestration
- eval.py → agents.py: runs · Reuses agent/evaluation logic
- agents.py → replay buffer: reads/writes · Collects and samples experience
- agents.py → env wrapper: interacts · Uses normalized environment APIs
- env wrapper → parallel rollout: uses · Executes rollouts in parallel
- parallel rollout → Atari env: wraps · Task-specific Atari environments
- parallel rollout → DMC env: wraps · Task-specific DMC environments
- parallel rollout → Memory maze: wraps · Task-specific memory-maze environments
- agents.py → world models.py: trains · Updates learned dynamics model
- agents.py → transformer model: trains · Alternative backbone path
- world models.py → attention blocks: uses · Relies on internal attention/sequence blocks
- world models.py → losses: optimizes with · Uses task and prediction losses
- world models.py → laprop: steps with · Uses optimization helpers
- transformer model → attention blocks: uses · Built from shared model blocks
- attention blocks → mamba models: implements · Sequence modules assembled into higher-level models
- mamba models → selective scan: calls · Invokes selective-scan operator boundary
- selective scan → selective scan cuda: dispatches · Calls compiled CUDA/C++ kernels
- mamba blocks → selective scan: depends on · Core modules use the operator boundary
- mamba models → triton ops: uses · Relies on Triton fused kernels for performance
- agents.py → tools.py: uses · Experiment helpers and logging utilities
- agents.py → utils.py: uses · Shared utility functions