Skip to content

Add seedability #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EthanAlmloff merged 3 commits into from
Oct 9, 2025
Merged

Add seedability #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cba9019
add seed ability to all 3 envs
EthanAlmloff Oct 9, 2025
6f131b1
add some additional rendering checks
EthanAlmloff Oct 9, 2025
9a43569
add determinism test to lunar lander
EthanAlmloff Oct 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
284 changes: 179 additions & 105 deletions src/box_2d/lunar_lander.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ use modurl::{
gym::{Gym, StepInfo},
spaces::{self, Space},
};
use rand::Rng;
use rand::distr::uniform::SampleRange;
use rand::distr::uniform::SampleUniform;
use std::cell::RefCell;
use std::rc::Rc;

Expand Down Expand Up @@ -54,6 +57,36 @@ const MAIN_ENGINE_Y_LOCATION: f32 = 4.0; // The Y location of the main engine on
const VIEWPORT_W: f32 = 600.0;
const VIEWPORT_H: f32 = 400.0;

enum EnvRng {
Thread(rand::rngs::ThreadRng),
Seeded(rand::rngs::StdRng),
}

impl EnvRng {
fn from_seed(seed: u64) -> Self {
EnvRng::Seeded(rand::SeedableRng::seed_from_u64(seed))
}
}

impl Default for EnvRng {
fn default() -> Self {
EnvRng::Thread(rand::rng())
}
}

impl EnvRng {
pub fn random_range<T, R>(&mut self, range: R) -> T
where
T: SampleUniform,
R: SampleRange<T>,
{
match self {
EnvRng::Thread(rng) => rng.random_range(range),
EnvRng::Seeded(rng) => rng.random_range(range),
}
}
}

// Rendering colors (ARGB format)
#[cfg(feature = "rendering")]
const COLOR_TERRAIN: u32 = 0xFFFFFFFF; // White
Expand Down Expand Up @@ -231,7 +264,7 @@ pub struct LunarLanderV3 {
torque_idx: i32,

// Random number generation
rng: rand::rngs::ThreadRng,
rng: EnvRng,

// Deterministic mode flag for testing
deterministic_mode: bool,
Expand All @@ -254,6 +287,7 @@ impl LunarLanderV3 {
#[cfg(feature = "rendering")]
#[builder(default = false)]
render: bool,
seed: Option<u64>,
) -> Self {
assert!(
-12.0 < gravity && gravity < 0.0,
Expand All @@ -275,6 +309,11 @@ impl LunarLanderV3 {
);
}

let mut rng = EnvRng::default();
if let Some(s) = seed {
rng = EnvRng::from_seed(s);
}

Self {
gravity,
enable_wind,
Expand All @@ -296,7 +335,7 @@ impl LunarLanderV3 {
sky_polys: Vec::new(),
wind_idx: 0,
torque_idx: 0,
rng: rand::rng(),
rng,
deterministic_mode: false,
device,
#[cfg(feature = "rendering")]
Expand All @@ -315,7 +354,7 @@ impl LunarLanderV3 {
fn destroy(&mut self) {
if let Some(world) = self.world.take() {
#[cfg(feature = "rendering")]
{
if self.renderer.is_some() {
// Clean up particles
for particle in &self.particles {
world.borrow_mut().destroy_body(particle.body.clone());
Expand Down Expand Up @@ -343,81 +382,6 @@ impl LunarLanderV3 {
}
}

#[allow(dead_code)]
fn get_raw_physics_state(&self) -> Result<Vec<f32>, candle_core::Error> {
if let Some(lander) = &self.lander {
let pos = lander.borrow().get_position();
let vel = lander.borrow().get_linear_velocity();
let angle = lander.borrow().get_angle();
let angular_vel = lander.borrow().get_angular_velocity();

// Get leg contacts (simplified)
let leg_contact_1 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[0])
.unwrap_or(false);
let leg_contact_2 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[1])
.unwrap_or(false);

// Return raw physics state (not normalized)
Ok(vec![
pos.x,
pos.y,
vel.x,
vel.y,
angle,
angular_vel,
if leg_contact_1 { 1.0 } else { 0.0 },
if leg_contact_2 { 1.0 } else { 0.0 },
])
} else {
Ok(vec![0.0f32; 8])
}
}

#[allow(dead_code)]
fn get_current_state(&self) -> Result<Tensor, candle_core::Error> {
if let Some(lander) = &self.lander {
let pos = lander.borrow().get_position();
let vel = lander.borrow().get_linear_velocity();
let angle = lander.borrow().get_angle();
let angular_vel = lander.borrow().get_angular_velocity();

// Get leg contacts (simplified - would need proper contact detection)
let leg_contact_1 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[0])
.unwrap_or(false);
let leg_contact_2 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[1])
.unwrap_or(false);

let state = vec![
(pos.x - VIEWPORT_W / SCALE / 2.0) / (VIEWPORT_W / SCALE / 2.0),
(pos.y - (self.helipad_y + LEG_DOWN / SCALE)) / (VIEWPORT_H / SCALE / 2.0),
vel.x * (VIEWPORT_W / SCALE / 2.0) / FPS,
vel.y * (VIEWPORT_H / SCALE / 2.0) / FPS,
angle,
20.0 * angular_vel / FPS,
if leg_contact_1 { 1.0 } else { 0.0 },
if leg_contact_2 { 1.0 } else { 0.0 },
];

Tensor::from_vec(state, vec![8], &self.device)
} else {
// Return zeros if no lander exists
let state = vec![0.0f32; 8];
Tensor::from_vec(state, vec![8], &self.device)
}
}

#[cfg(feature = "rendering")]
fn render(&mut self) {
if let Some(renderer) = &mut self.renderer
Expand Down Expand Up @@ -647,8 +611,6 @@ impl LunarLanderV3 {
tip: &(f32, f32),
dispersion: &[f32; 2],
) {
use rand::Rng;

let mut particle_body_def = B2bodyDef::default();
particle_body_def.body_type = B2bodyType::B2DynamicBody;

Expand Down Expand Up @@ -704,8 +666,6 @@ impl LunarLanderV3 {
direction: f32,
dispersion: &[f32; 2],
) {
use rand::Rng;

// Create particle for side engine flames (reduced count like original)
for _ in 0..1 {
let mut particle_body_def = B2bodyDef::default();
Expand Down Expand Up @@ -764,8 +724,6 @@ impl Gym for LunarLanderV3 {
type Error = candle_core::Error;

fn reset(&mut self) -> Result<Tensor, Self::Error> {
use rand::Rng;

self.destroy();

// Disable deterministic mode for normal reset
Expand Down Expand Up @@ -958,8 +916,6 @@ impl Gym for LunarLanderV3 {
}

fn step(&mut self, action: Tensor) -> Result<StepInfo, Self::Error> {
use rand::Rng;

assert!(self.lander.is_some(), "You forgot to call reset()");
let world = self.world.as_ref().unwrap().clone();
let lander = self.lander.as_ref().unwrap().clone();
Expand Down Expand Up @@ -1046,14 +1002,16 @@ impl Gym for LunarLanderV3 {
);

#[cfg(feature = "rendering")]
// Create main engine particles
self.create_main_engine_particles(
world.clone(),
&lander_pos,
lander_angle,
&tip,
&dispersion,
);
if self.renderer.is_some() {
// Create main engine particles
self.create_main_engine_particles(
world.clone(),
&lander_pos,
lander_angle,
&tip,
&dispersion,
);
}
}

let mut s_power = 0.0;
Expand Down Expand Up @@ -1089,23 +1047,25 @@ impl Gym for LunarLanderV3 {
);

#[cfg(feature = "rendering")]
// Create side engine particles
self.create_side_engine_particles(
world.clone(),
&lander_pos,
lander_angle,
&tip,
&side,
direction,
&dispersion,
);
if self.renderer.is_some() {
// Create side engine particles
self.create_side_engine_particles(
world.clone(),
&lander_pos,
lander_angle,
&tip,
&side,
direction,
&dispersion,
);
}
}

// Step the world
world.borrow_mut().step(1.0 / FPS, 6 * 30, 2 * 30);

#[cfg(feature = "rendering")]
{
if self.renderer.is_some() {
// Update and clean up particles
let dt = 1.0 / FPS;
self.particles.retain_mut(|particle| {
Expand Down Expand Up @@ -1244,6 +1204,46 @@ mod tests {
use super::*;
use crate::testing::{Testable, Tolerances, test_gym_against_python};

impl LunarLanderV3 {
fn get_current_state(&self) -> Result<Tensor, candle_core::Error> {
if let Some(lander) = &self.lander {
let pos = lander.borrow().get_position();
let vel = lander.borrow().get_linear_velocity();
let angle = lander.borrow().get_angle();
let angular_vel = lander.borrow().get_angular_velocity();

// Get leg contacts (simplified - would need proper contact detection)
let leg_contact_1 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[0])
.unwrap_or(false);
let leg_contact_2 = self
.contact_detector
.as_ref()
.map(|cd| cd.borrow().legs_ground_contact[1])
.unwrap_or(false);

let state = vec![
(pos.x - VIEWPORT_W / SCALE / 2.0) / (VIEWPORT_W / SCALE / 2.0),
(pos.y - (self.helipad_y + LEG_DOWN / SCALE)) / (VIEWPORT_H / SCALE / 2.0),
vel.x * (VIEWPORT_W / SCALE / 2.0) / FPS,
vel.y * (VIEWPORT_H / SCALE / 2.0) / FPS,
angle,
20.0 * angular_vel / FPS,
if leg_contact_1 { 1.0 } else { 0.0 },
if leg_contact_2 { 1.0 } else { 0.0 },
];

Tensor::from_vec(state, vec![8], &self.device)
} else {
// Return zeros if no lander exists
let state = vec![0.0f32; 8];
Tensor::from_vec(state, vec![8], &self.device)
}
}
}

impl Testable for LunarLanderV3 {
fn reset_deterministic(&mut self) -> Result<Tensor, candle_core::Error> {
// Do a deterministic reset that controls all random elements
Expand Down Expand Up @@ -1684,4 +1684,78 @@ mod tests {
}
}
}

#[test]
fn test_lunar_lander_determinism_seeded() {
let mut env1 = LunarLanderV3::builder().seed(42).enable_wind(true).build();
let mut env2 = LunarLanderV3::builder().seed(42).enable_wind(true).build();

let state1 = env1.reset_deterministic().expect("Failed to reset env1");
let state2 = env2.reset_deterministic().expect("Failed to reset env2");

// States should be the same size
assert_eq!(state1.dims(), state2.dims());
assert_eq!(state1.dims(), &[8]);

// The states might not be identical due to random terrain generation
// but they should be reasonable values
let state1_vec = state1
.to_vec1::<f32>()
.expect("Failed to convert state1 to vec");
let state2_vec = state2
.to_vec1::<f32>()
.expect("Failed to convert state2 to vec");

for (i, (&val1, &val2)) in state1_vec.iter().zip(state2_vec.iter()).enumerate() {
assert!(
val1.is_finite(),
"State {} should be finite, got {}",
i,
val1
);
assert!(
val2.is_finite(),
"State {} should be finite, got {}",
i,
val2
);
}

// Step both environments with the same actions and compare states
for step in 0..50 {
let action_value = step % 4; // Cycle through actions 0, 1, 2, 3
let action = Tensor::from_vec(vec![action_value as u32], vec![], &Device::Cpu)
.expect("Failed to create action tensor");

let step_info1 = env1.step(action.clone()).expect("Failed to step env1");
let step_info2 = env2.step(action).expect("Failed to step env2");

let state1 = step_info1.state;
let state2 = step_info2.state;

// States should be the same size
assert_eq!(state1.dims(), state2.dims());
assert_eq!(state1.dims(), &[8]);

let state1_vec = state1
.to_vec1::<f32>()
.expect("Failed to convert state1 to vec");
let state2_vec = state2
.to_vec1::<f32>()
.expect("Failed to convert state2 to vec");
for (i, (&val1, &val2)) in state1_vec.iter().zip(state2_vec.iter()).enumerate() {
assert!(
(val1 - val2).abs() < 1e-5,
"State {} differs between envs: {} vs {}",
i,
val1,
val2
);
}
if step_info1.done || step_info2.done {
env1.reset().expect("Failed to reset env1");
env2.reset().expect("Failed to reset env2");
}
}
}
}
Loading