huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix reinforcement learning example (#2837)

Browse Source
This commit is contained in:
LongYinan
2025-03-26 08:27:45 -07:00
committed by GitHub
parent 0d4097031c
commit cb02b389d5
3 changed files with 15 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
candle-examples/examples/reinforcement-learning/policy_gradient.rs
View File

@ -4,7 +4,7 @@ use candle_nn::{
linear, ops::log_softmax, ops::softmax, sequential::seq, Activation, AdamW, Optimizer,
ParamsAdamW, VarBuilder, VarMap,
};
use rand::{distributions::Distribution, rngs::ThreadRng, Rng};
use rand::{distr::Distribution, rngs::ThreadRng, Rng};
fn new_model(
input_shape: &[usize],
@ -39,7 +39,7 @@ fn accumulate_rewards(steps: &[Step<i64>]) -> Vec<f64> {
}
fn weighted_sample(probs: Vec<f32>, rng: &mut ThreadRng) -> Result<usize> {
let distribution = rand::distributions::WeightedIndex::new(probs).map_err(Error::wrap)?;
let distribution = rand::distr::weighted::WeightedIndex::new(probs).map_err(Error::wrap)?;
let mut rng = rng;
Ok(distribution.sample(&mut rng))
}
@ -65,10 +65,10 @@ pub fn run() -> Result<()> {
let mut optimizer = AdamW::new(varmap.all_vars(), optimizer_params)?;
let mut rng = rand::thread_rng();
let mut rng = rand::rng();
for epoch_idx in 0..100 {
let mut state = env.reset(rng.gen::<u64>())?;
let mut state = env.reset(rng.random::<u64>())?;
let mut steps: Vec<Step<i64>> = vec![];
loop {
@ -84,7 +84,7 @@ pub fn run() -> Result<()> {
steps.push(step.copy_with_obs(&state));
if step.terminated || step.truncated {
state = env.reset(rng.gen::<u64>())?;
state = env.reset(rng.random::<u64>())?;
if steps.len() > 5000 {
break;
}