rustytorch/crates/training/rtx-rl/README.md

RTX-RL: Reinforcement Learning at Scale

RTX-RL is a reinforcement learning crate designed for the RustyTorch++ ecosystem, implementing scalable RL algorithms with strict Test-Driven Development (TDD) principles.

Features

Core RL Algorithms

• PPO (Proximal Policy Optimization): On-policy algorithm with clipped objective
• SAC (Soft Actor-Critic): Off-policy algorithm with maximum entropy
• DPO (Direct Preference Optimization): Policy optimization from preference data

Infrastructure Components

• Experience Replay: Efficient circular buffer with prioritized sampling
• Environment Interface: Generic trait for RL environments
• Actor-Learner Architecture: Distributed RL training topology
• Standalone Implementation: Self-contained RL for testing and development

TDD Implementation

This crate demonstrates strict Test-Driven Development (TDD) following the Red-Green-Refactor cycle:

Phase 1: RED - Failing Tests First

1. Comprehensive test suite covering all RL components
2. Tests written before implementation
3. All tests initially fail (RED phase)

Phase 2: GREEN - Minimal Implementation

1. Minimal code to make tests pass
2. Focus on correctness over optimization
3. All tests pass (GREEN phase)

Phase 3: REFACTOR - Optimization

1. Improve performance and code quality
2. Maintain test coverage
3. Benchmark-driven optimization

Architecture

rtx-rl/
├── src/
│   ├── algorithms/          # RL algorithm implementations
│   │   ├── ppo.rs          # Proximal Policy Optimization
│   │   ├── sac.rs          # Soft Actor-Critic
│   │   └── dpo.rs          # Direct Preference Optimization
│   ├── environment.rs       # Environment trait and utilities
│   ├── replay_buffer.rs     # Experience storage and sampling
│   ├── actor_learner.rs     # Distributed training architecture
│   ├── standalone_rl.rs     # Self-contained RL implementation
│   └── tensor_utils.rs      # Tensor operation utilities
├── tests/                   # Comprehensive test coverage
├── benches/                 # Performance benchmarks
└── README.md               # This documentation

Usage

Basic Environment

use rtx_rl::{Environment, Step};

struct SimpleEnv {
    state: f32,
    steps: usize,
}

impl Environment for SimpleEnv {
    type State = f32;
    type Action = f32;
    type Reward = f32;

    fn reset(&mut self) -> Self::State {
        self.state = 0.0;
        self.steps = 0;
        self.state
    }

    fn step(&mut self, action: &Self::Action) -> Step<Self::State, Self::Reward> {
        self.state += action;
        self.steps += 1;
        
        Step {
            state: self.state,
            reward: -self.state.abs(), // Reward for staying near zero
            done: self.steps >= 100,
            info: std::collections::HashMap::new(),
        }
    }

    fn action_space(&self) -> (Self::Action, Self::Action) {
        (-1.0, 1.0)
    }

    fn observation_space(&self) -> (Self::State, Self::State) {
        (-10.0, 10.0)
    }
}

Standalone RL Training

use rtx_rl::standalone_rl::{PPO, ReplayBuffer, Experience};

fn main() {
    let mut ppo = PPO::new(4, 2); // 4D state, 2D action
    let mut buffer = ReplayBuffer::new(1000);
    
    // Training loop
    let mut state = vec![0.0; 4];
    
    for episode in 0..100 {
        for step in 0..50 {
            let action = ppo.get_action(&state);
            
            // Environment step (simplified)
            let next_state = vec![
                state[0] + action[0] * 0.1,
                state[1] + action[1] * 0.1,
                state[2] - state[0] * 0.05,
                state[3] - state[1] * 0.05,
            ];
            
            let reward = 1.0 - next_state.iter().map(|x| x.powi(2)).sum::<f32>().sqrt();
            
            let experience = Experience {
                state: state.clone(),
                action: action.clone(),
                reward,
                next_state: next_state.clone(),
                done: step >= 49,
            };
            
            buffer.push(experience);
            state = next_state;
        }
        
        // Train on collected experience
        if let Some(batch) = buffer.sample(32) {
            let loss = ppo.update(&batch);
            println!("Episode {}: Loss = {:.4}", episode, loss);
        }
    }
}

Testing

Run all tests:

cargo test

Run only library tests (standalone implementation):

cargo test --lib

Benchmarking

Run performance benchmarks:

cargo bench

Available benchmark suites:

• Matrix Operations: Basic linear algebra performance
• Replay Buffer: Experience storage and sampling throughput
• Policy Network: Neural network forward pass performance
• PPO Algorithm: Training algorithm performance
• Complete Episodes: End-to-end RL training performance
• Memory Usage: Memory allocation and management efficiency

Integration with RustyTorch++

RTX-RL is designed to integrate seamlessly with other RustyTorch++ crates:

• rtx-tensor: GPU-accelerated tensor operations
• rtx-runtime: CUDA memory management and kernel execution
• rtx-distributed: Multi-GPU and multi-node training
• rtx-graph: Computation graph optimization

Performance Goals

• Training Throughput: >10K samples/sec on single GPU
• Memory Efficiency: <100MB overhead for million-experience replay buffer
• Scaling: Linear scaling across multiple GPUs
• Latency: <1ms inference time for policy networks

Contributing

This crate follows strict TDD principles:

1. Write failing tests first (RED)
2. Implement minimal code to pass tests (GREEN)
3. Refactor for performance and maintainability (REFACTOR)
4. Maintain >90% test coverage
5. Include performance benchmarks for all major features

License

Part of the RustyTorch++ project. See workspace license for details.