201 lines
5.9 KiB
Markdown
201 lines
5.9 KiB
Markdown
# 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
|
|
```rust
|
|
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
|
|
```rust
|
|
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:
|
|
```bash
|
|
cargo test
|
|
```
|
|
|
|
Run only library tests (standalone implementation):
|
|
```bash
|
|
cargo test --lib
|
|
```
|
|
|
|
## Benchmarking
|
|
|
|
Run performance benchmarks:
|
|
```bash
|
|
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. |