475 lines
13 KiB
Rust
475 lines
13 KiB
Rust
//! Sample data and configurations for EmbodiedSim demo.
|
|
|
|
use embodied_shared::{
|
|
ActionDistribution, DomainRandomization, EnvironmentConfig, EnvironmentType, EvaluationRequest,
|
|
ImageEncoderType, ImaginationRequest, Observation, PolicyConfig, PolicyType, RewardType,
|
|
SafetyLimits, SimToRealConfig, TaskType, TrainingConfig, WorldModelConfig, WorldModelType,
|
|
};
|
|
|
|
// ============================================================================
|
|
// World Model Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a small RSSM configuration for fast training.
|
|
#[must_use]
|
|
pub fn small_rssm_config() -> WorldModelConfig {
|
|
WorldModelConfig {
|
|
model_type: WorldModelType::RSSM,
|
|
obs_embed_dim: 128,
|
|
action_embed_dim: 32,
|
|
deter_dim: 256,
|
|
stoch_dim: 16,
|
|
stoch_classes: 16,
|
|
hidden_dim: 256,
|
|
num_layers: 1,
|
|
image_encoder: ImageEncoderType::MLP,
|
|
predict_reward: true,
|
|
predict_done: true,
|
|
imagination_horizon: 10,
|
|
}
|
|
}
|
|
|
|
/// Create a medium RSSM configuration.
|
|
#[must_use]
|
|
pub fn medium_rssm_config() -> WorldModelConfig {
|
|
WorldModelConfig {
|
|
model_type: WorldModelType::RSSM,
|
|
obs_embed_dim: 256,
|
|
action_embed_dim: 64,
|
|
deter_dim: 512,
|
|
stoch_dim: 32,
|
|
stoch_classes: 32,
|
|
hidden_dim: 512,
|
|
num_layers: 2,
|
|
image_encoder: ImageEncoderType::CNN,
|
|
predict_reward: true,
|
|
predict_done: true,
|
|
imagination_horizon: 15,
|
|
}
|
|
}
|
|
|
|
/// Create a JEPA configuration.
|
|
#[must_use]
|
|
pub fn jepa_config() -> WorldModelConfig {
|
|
WorldModelConfig {
|
|
model_type: WorldModelType::JEPA,
|
|
obs_embed_dim: 256,
|
|
action_embed_dim: 64,
|
|
deter_dim: 512,
|
|
stoch_dim: 0,
|
|
stoch_classes: 0,
|
|
hidden_dim: 512,
|
|
num_layers: 4,
|
|
image_encoder: ImageEncoderType::ViT,
|
|
predict_reward: false,
|
|
predict_done: false,
|
|
imagination_horizon: 20,
|
|
}
|
|
}
|
|
|
|
/// Create a transformer world model configuration.
|
|
#[must_use]
|
|
pub fn transformer_wm_config() -> WorldModelConfig {
|
|
WorldModelConfig {
|
|
model_type: WorldModelType::Transformer,
|
|
obs_embed_dim: 256,
|
|
action_embed_dim: 64,
|
|
deter_dim: 512,
|
|
stoch_dim: 32,
|
|
stoch_classes: 32,
|
|
hidden_dim: 512,
|
|
num_layers: 6,
|
|
image_encoder: ImageEncoderType::ViT,
|
|
predict_reward: true,
|
|
predict_done: true,
|
|
imagination_horizon: 25,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Policy Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a small actor-critic policy.
|
|
#[must_use]
|
|
pub fn small_actor_critic_config() -> PolicyConfig {
|
|
PolicyConfig {
|
|
policy_type: PolicyType::ActorCritic,
|
|
state_dim: 256 + 16 * 16, // deter + stoch * classes
|
|
action_dim: 7,
|
|
actor_hidden: vec![128, 128],
|
|
critic_hidden: vec![128, 128],
|
|
gamma: 0.99,
|
|
gae_lambda: 0.95,
|
|
entropy_coef: 0.001,
|
|
continuous_actions: true,
|
|
action_dist: ActionDistribution::Normal,
|
|
}
|
|
}
|
|
|
|
/// Create a medium actor-critic policy.
|
|
#[must_use]
|
|
pub fn medium_actor_critic_config() -> PolicyConfig {
|
|
PolicyConfig {
|
|
policy_type: PolicyType::ActorCritic,
|
|
state_dim: 512 + 32 * 32,
|
|
action_dim: 7,
|
|
actor_hidden: vec![256, 256],
|
|
critic_hidden: vec![256, 256],
|
|
gamma: 0.997,
|
|
gae_lambda: 0.95,
|
|
entropy_coef: 3e-4,
|
|
continuous_actions: true,
|
|
action_dist: ActionDistribution::Normal,
|
|
}
|
|
}
|
|
|
|
/// Create an MPC policy configuration.
|
|
#[must_use]
|
|
pub fn mpc_policy_config() -> PolicyConfig {
|
|
PolicyConfig {
|
|
policy_type: PolicyType::MPC,
|
|
state_dim: 512,
|
|
action_dim: 7,
|
|
actor_hidden: vec![],
|
|
critic_hidden: vec![],
|
|
gamma: 0.99,
|
|
gae_lambda: 0.0,
|
|
entropy_coef: 0.0,
|
|
continuous_actions: true,
|
|
action_dist: ActionDistribution::Normal,
|
|
}
|
|
}
|
|
|
|
/// Create a diffusion policy configuration.
|
|
#[must_use]
|
|
pub fn diffusion_policy_config() -> PolicyConfig {
|
|
PolicyConfig {
|
|
policy_type: PolicyType::DiffusionPolicy,
|
|
state_dim: 512 + 32 * 32,
|
|
action_dim: 7,
|
|
actor_hidden: vec![256, 256, 256],
|
|
critic_hidden: vec![256, 256],
|
|
gamma: 0.99,
|
|
gae_lambda: 0.95,
|
|
entropy_coef: 0.0,
|
|
continuous_actions: true,
|
|
action_dist: ActionDistribution::Normal,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Environment Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a Panda reach task configuration.
|
|
#[must_use]
|
|
pub fn panda_reach_config() -> EnvironmentConfig {
|
|
EnvironmentConfig {
|
|
env_type: EnvironmentType::Panda,
|
|
task_type: TaskType::Reach,
|
|
num_joints: 7,
|
|
action_dim: 7,
|
|
obs_dim: 21,
|
|
use_camera: false,
|
|
camera_size: (64, 64),
|
|
max_episode_length: 100,
|
|
action_repeat: 2,
|
|
control_freq: 20.0,
|
|
reward_type: RewardType::Dense,
|
|
}
|
|
}
|
|
|
|
/// Create a Panda pick and place configuration.
|
|
#[must_use]
|
|
pub fn panda_pick_place_config() -> EnvironmentConfig {
|
|
EnvironmentConfig {
|
|
env_type: EnvironmentType::Panda,
|
|
task_type: TaskType::PickPlace,
|
|
num_joints: 7,
|
|
action_dim: 8, // 7 joints + gripper
|
|
obs_dim: 25, // joint state + object state
|
|
use_camera: true,
|
|
camera_size: (84, 84),
|
|
max_episode_length: 200,
|
|
action_repeat: 2,
|
|
control_freq: 20.0,
|
|
reward_type: RewardType::Shaped,
|
|
}
|
|
}
|
|
|
|
/// Create a UR5 push configuration.
|
|
#[must_use]
|
|
pub fn ur5_push_config() -> EnvironmentConfig {
|
|
EnvironmentConfig {
|
|
env_type: EnvironmentType::UR5,
|
|
task_type: TaskType::Push,
|
|
num_joints: 6,
|
|
action_dim: 6,
|
|
obs_dim: 18,
|
|
use_camera: true,
|
|
camera_size: (64, 64),
|
|
max_episode_length: 150,
|
|
action_repeat: 2,
|
|
control_freq: 20.0,
|
|
reward_type: RewardType::Dense,
|
|
}
|
|
}
|
|
|
|
/// Create a peg insertion configuration.
|
|
#[must_use]
|
|
pub fn peg_insertion_config() -> EnvironmentConfig {
|
|
EnvironmentConfig {
|
|
env_type: EnvironmentType::Panda,
|
|
task_type: TaskType::PegInsertion,
|
|
num_joints: 7,
|
|
action_dim: 8,
|
|
obs_dim: 28,
|
|
use_camera: true,
|
|
camera_size: (84, 84),
|
|
max_episode_length: 250,
|
|
action_repeat: 1,
|
|
control_freq: 50.0,
|
|
reward_type: RewardType::Shaped,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Training Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a fast training configuration for demos.
|
|
#[must_use]
|
|
pub fn fast_training_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
batch_size: 8,
|
|
sequence_length: 32,
|
|
wm_learning_rate: 1e-4,
|
|
actor_learning_rate: 3e-5,
|
|
critic_learning_rate: 3e-5,
|
|
epochs: 10,
|
|
buffer_size: 10_000,
|
|
prefill_steps: 500,
|
|
train_every: 5,
|
|
train_ratio: 64.0,
|
|
mixed_precision: false,
|
|
gradient_clip: 100.0,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
/// Create a standard training configuration.
|
|
#[must_use]
|
|
pub fn standard_training_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
batch_size: 64,
|
|
sequence_length: 64,
|
|
wm_learning_rate: 1e-4,
|
|
actor_learning_rate: 3e-5,
|
|
critic_learning_rate: 3e-5,
|
|
epochs: 100,
|
|
buffer_size: 1_000_000,
|
|
prefill_steps: 5000,
|
|
train_every: 5,
|
|
train_ratio: 512.0,
|
|
mixed_precision: true,
|
|
gradient_clip: 100.0,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
/// Create a DreamerV3-style training configuration.
|
|
#[must_use]
|
|
pub fn dreamerv3_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
batch_size: 16,
|
|
sequence_length: 64,
|
|
wm_learning_rate: 1e-4,
|
|
actor_learning_rate: 3e-5,
|
|
critic_learning_rate: 3e-5,
|
|
epochs: 1000,
|
|
buffer_size: 5_000_000,
|
|
prefill_steps: 2500,
|
|
train_every: 1,
|
|
train_ratio: 512.0,
|
|
mixed_precision: true,
|
|
gradient_clip: 1000.0,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Evaluation Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a quick evaluation request.
|
|
#[must_use]
|
|
pub fn quick_eval_request() -> EvaluationRequest {
|
|
EvaluationRequest {
|
|
num_episodes: 5,
|
|
render: false,
|
|
deterministic: true,
|
|
record_video: false,
|
|
env_config: panda_reach_config(),
|
|
}
|
|
}
|
|
|
|
/// Create a full evaluation request.
|
|
#[must_use]
|
|
pub fn full_eval_request() -> EvaluationRequest {
|
|
EvaluationRequest {
|
|
num_episodes: 100,
|
|
render: false,
|
|
deterministic: true,
|
|
record_video: false,
|
|
env_config: panda_reach_config(),
|
|
}
|
|
}
|
|
|
|
/// Create an evaluation with video recording.
|
|
#[must_use]
|
|
pub fn video_eval_request() -> EvaluationRequest {
|
|
EvaluationRequest {
|
|
num_episodes: 10,
|
|
render: true,
|
|
deterministic: true,
|
|
record_video: true,
|
|
env_config: panda_reach_config(),
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Imagination Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a short imagination request.
|
|
#[must_use]
|
|
pub fn short_imagination_request() -> ImaginationRequest {
|
|
ImaginationRequest {
|
|
initial_observation: Observation::default(),
|
|
horizon: 10,
|
|
num_trajectories: 5,
|
|
use_policy: true,
|
|
decode: false,
|
|
}
|
|
}
|
|
|
|
/// Create a long imagination request.
|
|
#[must_use]
|
|
pub fn long_imagination_request() -> ImaginationRequest {
|
|
ImaginationRequest {
|
|
initial_observation: Observation::default(),
|
|
horizon: 50,
|
|
num_trajectories: 20,
|
|
use_policy: true,
|
|
decode: true,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Sim-to-Real Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a conservative sim-to-real configuration.
|
|
#[must_use]
|
|
pub fn conservative_sim2real_config() -> SimToRealConfig {
|
|
SimToRealConfig {
|
|
domain_randomization: DomainRandomization {
|
|
dynamics: true,
|
|
visuals: true,
|
|
lighting: true,
|
|
mass_range: (0.9, 1.1),
|
|
friction_range: (0.7, 1.3),
|
|
damping_range: (0.9, 1.1),
|
|
action_noise: 0.005,
|
|
observation_noise: 0.005,
|
|
},
|
|
use_calibration: true,
|
|
action_scale: 0.8,
|
|
safety_limits: SafetyLimits {
|
|
max_joint_velocity: 1.0,
|
|
max_joint_acceleration: 3.0,
|
|
max_ee_velocity: 0.3,
|
|
max_force: 30.0,
|
|
workspace: [[-0.4, 0.4], [-0.4, 0.4], [0.1, 0.7]],
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create an aggressive sim-to-real configuration.
|
|
#[must_use]
|
|
pub fn aggressive_sim2real_config() -> SimToRealConfig {
|
|
SimToRealConfig {
|
|
domain_randomization: DomainRandomization {
|
|
dynamics: true,
|
|
visuals: true,
|
|
lighting: true,
|
|
mass_range: (0.5, 2.0),
|
|
friction_range: (0.2, 3.0),
|
|
damping_range: (0.5, 2.0),
|
|
action_noise: 0.02,
|
|
observation_noise: 0.02,
|
|
},
|
|
use_calibration: false,
|
|
action_scale: 1.0,
|
|
safety_limits: SafetyLimits::default(),
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_world_model_configs() {
|
|
let small = small_rssm_config();
|
|
let medium = medium_rssm_config();
|
|
|
|
assert!(small.hidden_dim < medium.hidden_dim);
|
|
assert!(small.deter_dim < medium.deter_dim);
|
|
}
|
|
|
|
#[test]
|
|
fn test_policy_configs() {
|
|
let small = small_actor_critic_config();
|
|
let medium = medium_actor_critic_config();
|
|
|
|
assert!(small.actor_hidden[0] < medium.actor_hidden[0]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_env_configs() {
|
|
let reach = panda_reach_config();
|
|
let pick = panda_pick_place_config();
|
|
|
|
assert_eq!(reach.task_type, TaskType::Reach);
|
|
assert_eq!(pick.task_type, TaskType::PickPlace);
|
|
assert!(pick.use_camera);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_configs() {
|
|
let fast = fast_training_config();
|
|
let standard = standard_training_config();
|
|
|
|
assert!(fast.epochs < standard.epochs);
|
|
assert!(fast.buffer_size < standard.buffer_size);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sim2real_configs() {
|
|
let conservative = conservative_sim2real_config();
|
|
let aggressive = aggressive_sim2real_config();
|
|
|
|
assert!(conservative.action_scale < aggressive.action_scale);
|
|
assert!(
|
|
conservative.domain_randomization.action_noise
|
|
< aggressive.domain_randomization.action_noise
|
|
);
|
|
}
|
|
}
|