Files
rustytorch/crates/training/rtx-rl/README.md
T
2026-03-04 00:08:42 +00:00

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.