316 lines
9.2 KiB
Rust
316 lines
9.2 KiB
Rust
//! Simplified robot simulator for the demo.
|
|
|
|
use embodied_shared::{Action, EnvironmentConfig, Observation, TaskType, TransitionInfo};
|
|
|
|
/// Robot simulator.
|
|
#[derive(Debug)]
|
|
pub struct Simulator {
|
|
/// Environment configuration.
|
|
config: EnvironmentConfig,
|
|
/// Current joint positions.
|
|
joint_positions: Vec<f32>,
|
|
/// Current joint velocities.
|
|
joint_velocities: Vec<f32>,
|
|
/// Target position.
|
|
target_position: [f32; 3],
|
|
/// Current timestep.
|
|
timestep: usize,
|
|
/// RNG state.
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl Simulator {
|
|
/// Create a new simulator.
|
|
pub fn new(config: &EnvironmentConfig) -> Self {
|
|
let mut sim = Self {
|
|
config: config.clone(),
|
|
joint_positions: vec![0.0; config.num_joints],
|
|
joint_velocities: vec![0.0; config.num_joints],
|
|
target_position: [0.3, 0.0, 0.5],
|
|
timestep: 0,
|
|
rng_state: 42,
|
|
};
|
|
sim.reset_state();
|
|
sim
|
|
}
|
|
|
|
/// Reset the environment.
|
|
pub fn reset(&mut self) -> Observation {
|
|
self.reset_state();
|
|
self.get_observation()
|
|
}
|
|
|
|
/// Reset internal state.
|
|
fn reset_state(&mut self) {
|
|
// Reset to home position
|
|
self.joint_positions = match self.config.num_joints {
|
|
7 => vec![0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785],
|
|
6 => vec![0.0, -1.57, 1.57, 0.0, 1.57, 0.0],
|
|
_ => vec![0.0; self.config.num_joints],
|
|
};
|
|
|
|
self.joint_velocities = vec![0.0; self.config.num_joints];
|
|
self.timestep = 0;
|
|
|
|
// Randomize target
|
|
self.target_position = [
|
|
0.3 + self.random() * 0.2 - 0.1,
|
|
0.0 + self.random() * 0.4 - 0.2,
|
|
0.5 + self.random() * 0.2 - 0.1,
|
|
];
|
|
}
|
|
|
|
/// Random number.
|
|
fn random(&mut self) -> f32 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f32 / (1u64 << 53) as f32
|
|
}
|
|
|
|
/// Step the environment.
|
|
pub fn step(&mut self, action: &Action) -> (Observation, f32, bool, TransitionInfo) {
|
|
self.timestep += 1;
|
|
|
|
// Apply action
|
|
if let Some(ref target_pos) = action.joint_positions {
|
|
for (i, &target) in target_pos.iter().enumerate() {
|
|
if i < self.joint_positions.len() {
|
|
// Simple P controller
|
|
let error = target - self.joint_positions[i];
|
|
self.joint_velocities[i] = error * 5.0; // Gain
|
|
self.joint_positions[i] += self.joint_velocities[i] * 0.02; // dt = 0.02s
|
|
|
|
// Clamp velocities
|
|
self.joint_velocities[i] = self.joint_velocities[i].clamp(-2.0, 2.0);
|
|
}
|
|
}
|
|
} else if let Some(ref target_vel) = action.joint_velocities {
|
|
for (i, &vel) in target_vel.iter().enumerate() {
|
|
if i < self.joint_velocities.len() {
|
|
self.joint_velocities[i] = vel.clamp(-2.0, 2.0);
|
|
self.joint_positions[i] += self.joint_velocities[i] * 0.02;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute end-effector position (simplified forward kinematics)
|
|
let ee_pos = self.compute_end_effector_position();
|
|
|
|
// Compute reward based on task
|
|
let (reward, success) = self.compute_reward(&ee_pos);
|
|
|
|
// Check if done
|
|
let done = success
|
|
|| self.timestep >= self.config.max_episode_length
|
|
|| self.is_out_of_bounds(&ee_pos);
|
|
|
|
let info = TransitionInfo {
|
|
success: Some(success),
|
|
distance_to_goal: Some(self.distance_to_target(&ee_pos)),
|
|
contact_forces: None,
|
|
};
|
|
|
|
(self.get_observation(), reward, done, info)
|
|
}
|
|
|
|
/// Get current observation.
|
|
fn get_observation(&self) -> Observation {
|
|
let ee_pos = self.compute_end_effector_position();
|
|
|
|
Observation {
|
|
joint_positions: self.joint_positions.clone(),
|
|
joint_velocities: self.joint_velocities.clone(),
|
|
end_effector_pose: [ee_pos[0], ee_pos[1], ee_pos[2], 1.0, 0.0, 0.0, 0.0],
|
|
camera_image: None,
|
|
image_dims: None,
|
|
timestamp: self.timestep as f64 * 0.02,
|
|
}
|
|
}
|
|
|
|
/// Compute end-effector position (simplified FK).
|
|
fn compute_end_effector_position(&self) -> [f32; 3] {
|
|
// Very simplified forward kinematics
|
|
let base_height = 0.3;
|
|
|
|
if self.joint_positions.len() >= 7 {
|
|
// Panda-like robot
|
|
let q = &self.joint_positions;
|
|
|
|
let x = 0.3 + 0.1 * q[0].sin() + 0.1 * q[1].sin() + 0.1 * q[3].sin();
|
|
let y = 0.1 * q[0].cos() * q[1].sin() + 0.1 * q[2].sin();
|
|
let z = base_height + 0.3 * q[1].cos() + 0.2 * q[3].cos() + 0.1 * q[5].cos();
|
|
|
|
[x, y, z]
|
|
} else {
|
|
// Generic
|
|
let x = 0.3 + 0.1 * self.joint_positions.iter().sum::<f32>();
|
|
let y = 0.0;
|
|
let z = base_height + 0.3;
|
|
|
|
[x, y, z]
|
|
}
|
|
}
|
|
|
|
/// Compute reward.
|
|
fn compute_reward(&self, ee_pos: &[f32; 3]) -> (f32, bool) {
|
|
let distance = self.distance_to_target(ee_pos);
|
|
let success = distance < 0.05; // 5cm threshold
|
|
|
|
match self.config.task_type {
|
|
TaskType::Reach => {
|
|
if success {
|
|
(10.0, true)
|
|
} else {
|
|
(-distance, false) // Dense reward
|
|
}
|
|
}
|
|
TaskType::Push | TaskType::PickPlace => {
|
|
// Simplified
|
|
if success {
|
|
(20.0, true)
|
|
} else {
|
|
(-distance * 2.0, false)
|
|
}
|
|
}
|
|
_ => {
|
|
// Default
|
|
if success {
|
|
(10.0, true)
|
|
} else {
|
|
(-distance, false)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Distance to target.
|
|
fn distance_to_target(&self, ee_pos: &[f32; 3]) -> f32 {
|
|
let dx = ee_pos[0] - self.target_position[0];
|
|
let dy = ee_pos[1] - self.target_position[1];
|
|
let dz = ee_pos[2] - self.target_position[2];
|
|
|
|
(dx * dx + dy * dy + dz * dz).sqrt()
|
|
}
|
|
|
|
/// Check if out of bounds.
|
|
fn is_out_of_bounds(&self, ee_pos: &[f32; 3]) -> bool {
|
|
ee_pos[0].abs() > 1.0 || ee_pos[1].abs() > 1.0 || ee_pos[2] < 0.0 || ee_pos[2] > 1.5
|
|
}
|
|
|
|
/// Get target position.
|
|
#[must_use]
|
|
pub fn target_position(&self) -> [f32; 3] {
|
|
self.target_position
|
|
}
|
|
|
|
/// Get current timestep.
|
|
#[must_use]
|
|
pub fn timestep(&self) -> usize {
|
|
self.timestep
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use embodied_shared::sample_env_config;
|
|
|
|
#[test]
|
|
fn test_simulator_creation() {
|
|
let config = sample_env_config();
|
|
let sim = Simulator::new(&config);
|
|
|
|
assert_eq!(sim.joint_positions.len(), config.num_joints);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reset() {
|
|
let config = sample_env_config();
|
|
let mut sim = Simulator::new(&config);
|
|
|
|
let obs = sim.reset();
|
|
assert_eq!(obs.joint_positions.len(), config.num_joints);
|
|
assert_eq!(sim.timestep(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_step() {
|
|
let config = sample_env_config();
|
|
let mut sim = Simulator::new(&config);
|
|
sim.reset();
|
|
|
|
let action = Action {
|
|
joint_positions: Some(vec![0.1; config.num_joints]),
|
|
..Default::default()
|
|
};
|
|
|
|
let (obs, reward, done, info) = sim.step(&action);
|
|
|
|
assert_eq!(obs.joint_positions.len(), config.num_joints);
|
|
assert!(reward.is_finite());
|
|
assert!(info.distance_to_goal.is_some());
|
|
assert_eq!(sim.timestep(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_steps() {
|
|
let config = sample_env_config();
|
|
let mut sim = Simulator::new(&config);
|
|
sim.reset();
|
|
|
|
for _ in 0..10 {
|
|
let action = Action {
|
|
joint_positions: Some(vec![0.1; config.num_joints]),
|
|
..Default::default()
|
|
};
|
|
let _ = sim.step(&action);
|
|
}
|
|
|
|
assert_eq!(sim.timestep(), 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reach_target() {
|
|
let config = sample_env_config();
|
|
let mut sim = Simulator::new(&config);
|
|
sim.reset();
|
|
|
|
// Force end-effector to target
|
|
sim.target_position = [0.3, 0.0, 0.5];
|
|
|
|
// Try to move toward target
|
|
for _ in 0..100 {
|
|
let action = Action {
|
|
joint_positions: Some(vec![0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785]),
|
|
..Default::default()
|
|
};
|
|
let (_, _, done, info) = sim.step(&action);
|
|
|
|
if done && info.success == Some(true) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_velocity_control() {
|
|
let config = sample_env_config();
|
|
let mut sim = Simulator::new(&config);
|
|
sim.reset();
|
|
|
|
let action = Action {
|
|
joint_velocities: Some(vec![0.1; config.num_joints]),
|
|
..Default::default()
|
|
};
|
|
|
|
let _ = sim.step(&action);
|
|
|
|
// Velocities should be applied
|
|
for &vel in &sim.joint_velocities {
|
|
assert!(vel.abs() <= 2.0);
|
|
}
|
|
}
|
|
}
|