302 lines
9.7 KiB
Rust
302 lines
9.7 KiB
Rust
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
|
|
use rtx_rl::standalone_rl::{Experience, Matrix, PPO, ReplayBuffer};
|
|
|
|
fn bench_matrix_operations(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("matrix_operations");
|
|
|
|
for size in [10, 100, 1000].iter() {
|
|
group.throughput(Throughput::Elements(*size as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("matrix_creation", size),
|
|
size,
|
|
|b, &size| {
|
|
b.iter(|| Matrix::randn(vec![size, size]));
|
|
},
|
|
);
|
|
|
|
let matrix_a = Matrix::randn(vec![*size, *size]);
|
|
let matrix_b = Matrix::randn(vec![*size, *size]);
|
|
|
|
group.bench_with_input(BenchmarkId::new("matrix_addition", size), size, |b, _| {
|
|
b.iter(|| matrix_a.add(&matrix_b));
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("matrix_scalar_mul", size), size, |b, _| {
|
|
b.iter(|| matrix_a.mul(0.5));
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("matrix_mean", size), size, |b, _| {
|
|
b.iter(|| matrix_a.mean());
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_replay_buffer_operations(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("replay_buffer");
|
|
|
|
for capacity in [1000, 10000, 100000].iter() {
|
|
group.throughput(Throughput::Elements(*capacity as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("buffer_fill", capacity),
|
|
capacity,
|
|
|b, &cap| {
|
|
b.iter(|| {
|
|
let mut buffer = ReplayBuffer::new(cap);
|
|
for i in 0..cap {
|
|
let exp = Experience {
|
|
state: vec![i as f32, (i * 2) as f32],
|
|
action: vec![(i % 3) as f32],
|
|
reward: (i as f32) * 0.01,
|
|
next_state: vec![(i + 1) as f32, ((i + 1) * 2) as f32],
|
|
done: i == cap - 1,
|
|
};
|
|
buffer.push(exp);
|
|
}
|
|
});
|
|
},
|
|
);
|
|
|
|
let mut buffer = ReplayBuffer::new(*capacity);
|
|
for i in 0..*capacity {
|
|
let exp = Experience {
|
|
state: vec![i as f32, (i * 2) as f32],
|
|
action: vec![(i % 3) as f32],
|
|
reward: (i as f32) * 0.01,
|
|
next_state: vec![(i + 1) as f32, ((i + 1) * 2) as f32],
|
|
done: false,
|
|
};
|
|
buffer.push(exp);
|
|
}
|
|
|
|
for batch_size in [32, 128, 512].iter() {
|
|
if *batch_size <= *capacity {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("buffer_sample", format!("{}_{}", capacity, batch_size)),
|
|
&(*capacity, *batch_size),
|
|
|b, &(_cap, batch)| {
|
|
b.iter(|| buffer.sample(batch).unwrap());
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_policy_network(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("policy_network");
|
|
|
|
for (input_dim, output_dim) in [(4, 2), (16, 8), (64, 32), (256, 128)].iter() {
|
|
group.throughput(Throughput::Elements(*input_dim as u64));
|
|
|
|
let network = rtx_rl::standalone_rl::PolicyNetwork::new(*input_dim, *output_dim);
|
|
let input: Vec<f32> = (0..*input_dim).map(|i| (i as f32) * 0.1).collect();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("forward_pass", format!("{}x{}", input_dim, output_dim)),
|
|
&(*input_dim, *output_dim),
|
|
|b, _| {
|
|
b.iter(|| network.forward(&input));
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_ppo_algorithm(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("ppo");
|
|
|
|
for batch_size in [32, 128, 512].iter() {
|
|
group.throughput(Throughput::Elements(*batch_size as u64));
|
|
|
|
let mut ppo = PPO::new(8, 4);
|
|
let mut experiences = Vec::new();
|
|
|
|
for i in 0..*batch_size {
|
|
let exp = Experience {
|
|
state: vec![
|
|
(i as f32) * 0.01,
|
|
(i as f32) * 0.02,
|
|
(i as f32) * 0.03,
|
|
(i as f32) * 0.04,
|
|
(i as f32) * 0.05,
|
|
(i as f32) * 0.06,
|
|
(i as f32) * 0.07,
|
|
(i as f32) * 0.08,
|
|
],
|
|
action: vec![
|
|
((i % 4) as f32) * 0.25,
|
|
((i % 3) as f32) * 0.33,
|
|
((i % 2) as f32) * 0.5,
|
|
(i as f32) * 0.01,
|
|
],
|
|
reward: if i % 2 == 0 { 1.0 } else { -0.5 },
|
|
next_state: vec![
|
|
(i as f32 + 1.0) * 0.01,
|
|
(i as f32 + 1.0) * 0.02,
|
|
(i as f32 + 1.0) * 0.03,
|
|
(i as f32 + 1.0) * 0.04,
|
|
(i as f32 + 1.0) * 0.05,
|
|
(i as f32 + 1.0) * 0.06,
|
|
(i as f32 + 1.0) * 0.07,
|
|
(i as f32 + 1.0) * 0.08,
|
|
],
|
|
done: i == *batch_size - 1,
|
|
};
|
|
experiences.push(exp);
|
|
}
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("get_action", batch_size),
|
|
batch_size,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
for exp in &experiences {
|
|
ppo.get_action(&exp.state);
|
|
}
|
|
});
|
|
},
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("get_value", batch_size),
|
|
batch_size,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
for exp in &experiences {
|
|
ppo.get_value(&exp.state);
|
|
}
|
|
});
|
|
},
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("compute_advantages", batch_size),
|
|
batch_size,
|
|
|b, _| {
|
|
b.iter(|| ppo.compute_advantages(&experiences));
|
|
},
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("update", batch_size),
|
|
batch_size,
|
|
|b, _| {
|
|
b.iter(|| ppo.update(&experiences));
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_complete_rl_episode(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("complete_rl_episode");
|
|
|
|
for episode_length in [10, 100, 1000].iter() {
|
|
group.throughput(Throughput::Elements(*episode_length as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("full_episode", episode_length),
|
|
episode_length,
|
|
|b, &length| {
|
|
b.iter(|| {
|
|
let mut ppo = PPO::new(4, 2);
|
|
let mut buffer = ReplayBuffer::new(length * 2);
|
|
|
|
// Simulate episode
|
|
let mut state = vec![0.0, 0.0, 0.0, 0.0];
|
|
let mut total_reward = 0.0;
|
|
|
|
for step in 0..length {
|
|
let action = ppo.get_action(&state);
|
|
|
|
// Simple environment dynamics
|
|
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 distance = next_state.iter().map(|x| x.powi(2)).sum::<f32>().sqrt();
|
|
let reward = 1.0 - distance.min(1.0);
|
|
total_reward += reward;
|
|
|
|
let experience = Experience {
|
|
state: state.clone(),
|
|
action: action.clone(),
|
|
reward,
|
|
next_state: next_state.clone(),
|
|
done: step >= length - 1,
|
|
};
|
|
|
|
buffer.push(experience);
|
|
state = next_state;
|
|
|
|
// Train every 32 steps
|
|
if step % 32 == 31 {
|
|
if let Some(batch) = buffer.sample(32.min(buffer.len())) {
|
|
ppo.update(&batch);
|
|
}
|
|
}
|
|
}
|
|
|
|
total_reward
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_memory_usage(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("memory_usage");
|
|
|
|
for num_experiences in [1000, 10000, 100000].iter() {
|
|
group.throughput(Throughput::Elements(*num_experiences as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("experience_creation", num_experiences),
|
|
num_experiences,
|
|
|b, &count| {
|
|
b.iter(|| {
|
|
let mut experiences = Vec::with_capacity(count);
|
|
for i in 0..count {
|
|
let exp = Experience {
|
|
state: vec![i as f32; 8],
|
|
action: vec![(i % 4) as f32; 4],
|
|
reward: (i as f32) * 0.001,
|
|
next_state: vec![(i + 1) as f32; 8],
|
|
done: i == count - 1,
|
|
};
|
|
experiences.push(exp);
|
|
}
|
|
experiences
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_matrix_operations,
|
|
bench_replay_buffer_operations,
|
|
bench_policy_network,
|
|
bench_ppo_algorithm,
|
|
bench_complete_rl_episode,
|
|
bench_memory_usage
|
|
);
|
|
|
|
criterion_main!(benches);
|