591 lines
18 KiB
Rust
591 lines
18 KiB
Rust
//! EmbodiedSim - Robotics World Model Trainer.
|
|
//!
|
|
//! This demo showcases world models for sim-to-real robotics transfer,
|
|
//! implementing RSSM and JEPA architectures for learning robot dynamics.
|
|
|
|
pub mod actor_critic;
|
|
pub mod jepa;
|
|
pub mod replay_buffer;
|
|
pub mod rssm;
|
|
pub mod sample_data;
|
|
pub mod simulator;
|
|
|
|
use thiserror::Error;
|
|
|
|
use embodied_shared::{
|
|
Action, EnvironmentConfig, EvaluationRequest, EvaluationResult, ImaginationRequest,
|
|
ImaginedTrajectory, Observation, PolicyConfig, TrainingConfig, TrainingProgress, Transition,
|
|
WorldModelConfig,
|
|
};
|
|
|
|
/// Errors that can occur in EmbodiedSim.
|
|
#[derive(Debug, Error)]
|
|
pub enum EmbodiedError {
|
|
/// Invalid configuration.
|
|
#[error("Invalid configuration: {0}")]
|
|
InvalidConfig(String),
|
|
|
|
/// Training failed.
|
|
#[error("Training failed: {0}")]
|
|
TrainingFailed(String),
|
|
|
|
/// World model not initialized.
|
|
#[error("World model not initialized")]
|
|
ModelNotInitialized,
|
|
|
|
/// Policy not initialized.
|
|
#[error("Policy not initialized")]
|
|
PolicyNotInitialized,
|
|
|
|
/// Buffer empty.
|
|
#[error("Replay buffer is empty")]
|
|
BufferEmpty,
|
|
|
|
/// Imagination failed.
|
|
#[error("Imagination failed: {0}")]
|
|
ImaginationFailed(String),
|
|
}
|
|
|
|
/// Main EmbodiedSim system.
|
|
#[derive(Debug)]
|
|
pub struct EmbodiedSim {
|
|
/// World model configuration.
|
|
wm_config: WorldModelConfig,
|
|
/// Policy configuration.
|
|
policy_config: PolicyConfig,
|
|
/// Environment configuration.
|
|
env_config: EnvironmentConfig,
|
|
/// RSSM world model.
|
|
rssm: rssm::RSSM,
|
|
/// Actor-Critic policy.
|
|
policy: actor_critic::ActorCritic,
|
|
/// Replay buffer.
|
|
buffer: replay_buffer::ReplayBuffer,
|
|
/// Simulator.
|
|
simulator: simulator::Simulator,
|
|
/// Is initialized.
|
|
initialized: bool,
|
|
}
|
|
|
|
impl Default for EmbodiedSim {
|
|
fn default() -> Self {
|
|
Self::new(
|
|
WorldModelConfig::default(),
|
|
PolicyConfig::default(),
|
|
EnvironmentConfig::default(),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl EmbodiedSim {
|
|
/// Create a new EmbodiedSim instance.
|
|
#[must_use]
|
|
pub fn new(
|
|
wm_config: WorldModelConfig,
|
|
policy_config: PolicyConfig,
|
|
env_config: EnvironmentConfig,
|
|
) -> Self {
|
|
Self {
|
|
rssm: rssm::RSSM::new(&wm_config),
|
|
policy: actor_critic::ActorCritic::new(&policy_config),
|
|
buffer: replay_buffer::ReplayBuffer::new(100_000),
|
|
simulator: simulator::Simulator::new(&env_config),
|
|
wm_config,
|
|
policy_config,
|
|
env_config,
|
|
initialized: true,
|
|
}
|
|
}
|
|
|
|
/// Add transition to replay buffer.
|
|
pub fn add_transition(&mut self, transition: Transition) {
|
|
self.buffer.add(transition);
|
|
}
|
|
|
|
/// Collect experience from environment.
|
|
pub fn collect_experience(
|
|
&mut self,
|
|
num_steps: usize,
|
|
) -> Result<Vec<Transition>, EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::ModelNotInitialized);
|
|
}
|
|
|
|
let mut transitions = Vec::with_capacity(num_steps);
|
|
let mut obs = self.simulator.reset();
|
|
|
|
for _ in 0..num_steps {
|
|
// Get action from policy
|
|
let state = self.rssm.encode_observation(&obs);
|
|
let action = self.policy.sample_action(&state);
|
|
let action_struct = self.action_from_vec(&action);
|
|
|
|
// Step environment
|
|
let (next_obs, reward, done, info) = self.simulator.step(&action_struct);
|
|
|
|
let transition = Transition {
|
|
observation: obs.clone(),
|
|
action: action_struct,
|
|
reward,
|
|
next_observation: next_obs.clone(),
|
|
done,
|
|
truncated: false,
|
|
info,
|
|
};
|
|
|
|
transitions.push(transition.clone());
|
|
self.buffer.add(transition);
|
|
|
|
if done {
|
|
obs = self.simulator.reset();
|
|
} else {
|
|
obs = next_obs;
|
|
}
|
|
}
|
|
|
|
Ok(transitions)
|
|
}
|
|
|
|
/// Convert action vector to Action struct.
|
|
fn action_from_vec(&self, action: &[f64]) -> Action {
|
|
Action {
|
|
joint_positions: Some(action.iter().map(|&x| x as f32).collect()),
|
|
joint_velocities: None,
|
|
joint_torques: None,
|
|
ee_delta: None,
|
|
gripper: if action.len() > self.env_config.num_joints {
|
|
Some(action[self.env_config.num_joints] as f32)
|
|
} else {
|
|
None
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Train the world model.
|
|
pub fn train_world_model(
|
|
&mut self,
|
|
config: &TrainingConfig,
|
|
progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
|
|
) -> Result<(), EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::ModelNotInitialized);
|
|
}
|
|
|
|
if self.buffer.len() < config.batch_size * config.sequence_length {
|
|
return Err(EmbodiedError::BufferEmpty);
|
|
}
|
|
|
|
let start_time = std::time::Instant::now();
|
|
let total_steps = config.epochs * 100; // 100 batches per epoch
|
|
|
|
for epoch in 0..config.epochs {
|
|
for batch_idx in 0..100 {
|
|
// Sample batch from buffer
|
|
let batch = self
|
|
.buffer
|
|
.sample_sequence(config.batch_size, config.sequence_length);
|
|
|
|
// Train world model on batch
|
|
let loss = self.rssm.train_step(&batch, config.wm_learning_rate);
|
|
|
|
let step = epoch * 100 + batch_idx + 1;
|
|
|
|
if let Some(ref callback) = progress_callback {
|
|
callback(TrainingProgress {
|
|
epoch: epoch + 1,
|
|
total_epochs: config.epochs,
|
|
step,
|
|
total_steps,
|
|
wm_loss: loss,
|
|
image_loss: Some(loss * 0.6),
|
|
reward_loss: Some(loss * 0.2),
|
|
kl_loss: Some(loss * 0.2),
|
|
actor_loss: None,
|
|
critic_loss: None,
|
|
entropy: None,
|
|
learning_rate: config.wm_learning_rate,
|
|
env_return: None,
|
|
episode_length: None,
|
|
elapsed_seconds: start_time.elapsed().as_secs_f64(),
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Train the policy using imagined rollouts.
|
|
pub fn train_policy(
|
|
&mut self,
|
|
config: &TrainingConfig,
|
|
progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
|
|
) -> Result<(), EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::ModelNotInitialized);
|
|
}
|
|
|
|
let start_time = std::time::Instant::now();
|
|
let total_steps = config.epochs * 50; // 50 updates per epoch
|
|
|
|
for epoch in 0..config.epochs {
|
|
for update_idx in 0..50 {
|
|
// Sample starting states from buffer
|
|
let batch = self.buffer.sample(config.batch_size);
|
|
|
|
// Imagine trajectories
|
|
let trajectories = self.imagine_batch(&batch, self.wm_config.imagination_horizon);
|
|
|
|
// Compute returns and advantages
|
|
let (returns, advantages) = self.compute_advantages(&trajectories);
|
|
|
|
// Update policy
|
|
let (actor_loss, critic_loss, entropy) =
|
|
self.policy
|
|
.update(&trajectories, &returns, &advantages, config);
|
|
|
|
let step = epoch * 50 + update_idx + 1;
|
|
|
|
if let Some(ref callback) = progress_callback {
|
|
callback(TrainingProgress {
|
|
epoch: epoch + 1,
|
|
total_epochs: config.epochs,
|
|
step,
|
|
total_steps,
|
|
wm_loss: 0.0,
|
|
image_loss: None,
|
|
reward_loss: None,
|
|
kl_loss: None,
|
|
actor_loss: Some(actor_loss),
|
|
critic_loss: Some(critic_loss),
|
|
entropy: Some(entropy),
|
|
learning_rate: config.actor_learning_rate,
|
|
env_return: None,
|
|
episode_length: None,
|
|
elapsed_seconds: start_time.elapsed().as_secs_f64(),
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Imagine trajectories from a batch of observations.
|
|
fn imagine_batch(&self, batch: &[Transition], horizon: usize) -> Vec<ImaginedTrajectory> {
|
|
batch
|
|
.iter()
|
|
.map(|trans| self.imagine_from_obs(&trans.observation, horizon))
|
|
.collect()
|
|
}
|
|
|
|
/// Imagine a trajectory from an observation.
|
|
fn imagine_from_obs(&self, obs: &Observation, horizon: usize) -> ImaginedTrajectory {
|
|
let mut states = Vec::with_capacity(horizon);
|
|
let mut actions = Vec::with_capacity(horizon);
|
|
let mut rewards = Vec::with_capacity(horizon);
|
|
let mut values = Vec::with_capacity(horizon);
|
|
let mut dones = Vec::with_capacity(horizon);
|
|
|
|
let mut state = self.rssm.encode_observation(obs);
|
|
|
|
for _ in 0..horizon {
|
|
// Sample action from policy
|
|
let action = self.policy.sample_action(&state);
|
|
|
|
// Predict value
|
|
let value = self.policy.predict_value(&state);
|
|
|
|
// Imagine next state
|
|
let (next_state, reward, done) = self.rssm.imagine_step(&state, &action);
|
|
|
|
states.push(state.clone());
|
|
actions.push(action);
|
|
rewards.push(reward);
|
|
values.push(value);
|
|
dones.push(done);
|
|
|
|
state = next_state;
|
|
}
|
|
|
|
ImaginedTrajectory {
|
|
states,
|
|
actions,
|
|
rewards,
|
|
values,
|
|
dones,
|
|
decoded_observations: None,
|
|
}
|
|
}
|
|
|
|
/// Compute returns and advantages for imagined trajectories.
|
|
fn compute_advantages(
|
|
&self,
|
|
trajectories: &[ImaginedTrajectory],
|
|
) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
|
|
let gamma = self.policy_config.gamma as f64;
|
|
let lambda = self.policy_config.gae_lambda as f64;
|
|
|
|
let mut all_returns = Vec::with_capacity(trajectories.len());
|
|
let mut all_advantages = Vec::with_capacity(trajectories.len());
|
|
|
|
for traj in trajectories {
|
|
let horizon = traj.rewards.len();
|
|
let mut returns = vec![0.0; horizon];
|
|
let mut advantages = vec![0.0; horizon];
|
|
|
|
// Bootstrap from last value
|
|
let mut gae = 0.0;
|
|
|
|
for t in (0..horizon).rev() {
|
|
let next_value = if t + 1 < horizon {
|
|
traj.values[t + 1]
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let delta =
|
|
traj.rewards[t] + gamma * next_value * (1.0 - traj.dones[t]) - traj.values[t];
|
|
gae = delta + gamma * lambda * (1.0 - traj.dones[t]) * gae;
|
|
advantages[t] = gae;
|
|
returns[t] = advantages[t] + traj.values[t];
|
|
}
|
|
|
|
all_returns.push(returns);
|
|
all_advantages.push(advantages);
|
|
}
|
|
|
|
(all_returns, all_advantages)
|
|
}
|
|
|
|
/// Imagine future trajectories.
|
|
pub fn imagine(
|
|
&self,
|
|
request: &ImaginationRequest,
|
|
) -> Result<Vec<ImaginedTrajectory>, EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::ModelNotInitialized);
|
|
}
|
|
|
|
let mut trajectories = Vec::with_capacity(request.num_trajectories);
|
|
|
|
for _ in 0..request.num_trajectories {
|
|
let traj = self.imagine_from_obs(&request.initial_observation, request.horizon);
|
|
trajectories.push(traj);
|
|
}
|
|
|
|
Ok(trajectories)
|
|
}
|
|
|
|
/// Evaluate the policy.
|
|
pub fn evaluate(
|
|
&mut self,
|
|
request: &EvaluationRequest,
|
|
) -> Result<EvaluationResult, EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::ModelNotInitialized);
|
|
}
|
|
|
|
let start_time = std::time::Instant::now();
|
|
let mut episode_returns = Vec::with_capacity(request.num_episodes);
|
|
let mut episode_lengths = Vec::with_capacity(request.num_episodes);
|
|
let mut successes = 0;
|
|
|
|
for _ in 0..request.num_episodes {
|
|
let mut obs = self.simulator.reset();
|
|
let mut episode_return = 0.0;
|
|
let mut episode_length = 0;
|
|
|
|
loop {
|
|
let state = self.rssm.encode_observation(&obs);
|
|
let action = if request.deterministic {
|
|
self.policy.get_action_mean(&state)
|
|
} else {
|
|
self.policy.sample_action(&state)
|
|
};
|
|
|
|
let action_struct = self.action_from_vec(&action);
|
|
let (next_obs, reward, done, info) = self.simulator.step(&action_struct);
|
|
|
|
episode_return += reward as f64;
|
|
episode_length += 1;
|
|
|
|
if info.success == Some(true) {
|
|
successes += 1;
|
|
}
|
|
|
|
if done || episode_length >= self.env_config.max_episode_length {
|
|
break;
|
|
}
|
|
|
|
obs = next_obs;
|
|
}
|
|
|
|
episode_returns.push(episode_return);
|
|
episode_lengths.push(episode_length);
|
|
}
|
|
|
|
let mean_return = episode_returns.iter().sum::<f64>() / request.num_episodes as f64;
|
|
let mean_length =
|
|
episode_lengths.iter().sum::<usize>() as f64 / request.num_episodes as f64;
|
|
|
|
let std_return = if request.num_episodes > 1 {
|
|
let variance = episode_returns
|
|
.iter()
|
|
.map(|r| (r - mean_return).powi(2))
|
|
.sum::<f64>()
|
|
/ (request.num_episodes - 1) as f64;
|
|
variance.sqrt()
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
Ok(EvaluationResult {
|
|
mean_return,
|
|
std_return,
|
|
mean_length,
|
|
success_rate: successes as f64 / request.num_episodes as f64,
|
|
episode_returns,
|
|
episode_lengths,
|
|
eval_time: start_time.elapsed().as_secs_f64(),
|
|
})
|
|
}
|
|
|
|
/// Get action for a given observation.
|
|
pub fn get_action(
|
|
&self,
|
|
obs: &Observation,
|
|
deterministic: bool,
|
|
) -> Result<Action, EmbodiedError> {
|
|
if !self.initialized {
|
|
return Err(EmbodiedError::PolicyNotInitialized);
|
|
}
|
|
|
|
let state = self.rssm.encode_observation(obs);
|
|
let action = if deterministic {
|
|
self.policy.get_action_mean(&state)
|
|
} else {
|
|
self.policy.sample_action(&state)
|
|
};
|
|
|
|
Ok(self.action_from_vec(&action))
|
|
}
|
|
|
|
/// Get world model config.
|
|
#[must_use]
|
|
pub fn world_model_config(&self) -> &WorldModelConfig {
|
|
&self.wm_config
|
|
}
|
|
|
|
/// Get policy config.
|
|
#[must_use]
|
|
pub fn policy_config(&self) -> &PolicyConfig {
|
|
&self.policy_config
|
|
}
|
|
|
|
/// Get environment config.
|
|
#[must_use]
|
|
pub fn env_config(&self) -> &EnvironmentConfig {
|
|
&self.env_config
|
|
}
|
|
|
|
/// Get buffer size.
|
|
#[must_use]
|
|
pub fn buffer_size(&self) -> usize {
|
|
self.buffer.len()
|
|
}
|
|
}
|
|
|
|
/// Run the demo.
|
|
pub fn run_demo() -> Result<EvaluationResult, EmbodiedError> {
|
|
let mut embodied = EmbodiedSim::default();
|
|
|
|
// Collect some experience
|
|
let _ = embodied.collect_experience(100)?;
|
|
|
|
// Evaluate
|
|
let request = EvaluationRequest {
|
|
num_episodes: 5,
|
|
..Default::default()
|
|
};
|
|
embodied.evaluate(&request)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_embodied_creation() {
|
|
let embodied = EmbodiedSim::default();
|
|
assert!(embodied.initialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_collect_experience() {
|
|
let mut embodied = EmbodiedSim::default();
|
|
let result = embodied.collect_experience(10);
|
|
assert!(result.is_ok());
|
|
|
|
let transitions = result.unwrap();
|
|
assert_eq!(transitions.len(), 10);
|
|
assert_eq!(embodied.buffer_size(), 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_action() {
|
|
let embodied = EmbodiedSim::default();
|
|
let obs = embodied_shared::sample_observation();
|
|
|
|
let action = embodied.get_action(&obs, true);
|
|
assert!(action.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_imagine() {
|
|
let mut embodied = EmbodiedSim::default();
|
|
let _ = embodied.collect_experience(10);
|
|
|
|
let request = ImaginationRequest {
|
|
initial_observation: embodied_shared::sample_observation(),
|
|
horizon: 5,
|
|
num_trajectories: 3,
|
|
use_policy: true,
|
|
decode: false,
|
|
};
|
|
|
|
let result = embodied.imagine(&request);
|
|
assert!(result.is_ok());
|
|
|
|
let trajectories = result.unwrap();
|
|
assert_eq!(trajectories.len(), 3);
|
|
for traj in &trajectories {
|
|
assert_eq!(traj.states.len(), 5);
|
|
assert_eq!(traj.rewards.len(), 5);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_evaluate() {
|
|
let mut embodied = EmbodiedSim::default();
|
|
let _ = embodied.collect_experience(10);
|
|
|
|
let request = EvaluationRequest {
|
|
num_episodes: 2,
|
|
..Default::default()
|
|
};
|
|
|
|
let result = embodied.evaluate(&request);
|
|
assert!(result.is_ok());
|
|
|
|
let eval = result.unwrap();
|
|
assert_eq!(eval.episode_returns.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_demo() {
|
|
let result = run_demo();
|
|
assert!(result.is_ok());
|
|
}
|
|
}
|