//! Shared types for EmbodiedSim - Robotics World Model Trainer. //! //! This crate defines the IPC types for sim-to-real robotics world models. use serde::{Deserialize, Serialize}; // ============================================================================ // Observation Types // ============================================================================ /// Robot observation (state). #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Observation { /// Joint positions (radians). pub joint_positions: Vec, /// Joint velocities (rad/s). pub joint_velocities: Vec, /// End-effector pose (x, y, z, qw, qx, qy, qz). pub end_effector_pose: [f32; 7], /// Camera image (flattened RGB). pub camera_image: Option>, /// Image dimensions (width, height). pub image_dims: Option<(usize, usize)>, /// Timestamp. pub timestamp: f64, } impl Default for Observation { fn default() -> Self { Self { joint_positions: vec![0.0; 7], joint_velocities: vec![0.0; 7], end_effector_pose: [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], camera_image: None, image_dims: None, timestamp: 0.0, } } } /// Action command to robot. #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Action { /// Joint position targets (radians). pub joint_positions: Option>, /// Joint velocity targets (rad/s). pub joint_velocities: Option>, /// Joint torques (Nm). pub joint_torques: Option>, /// End-effector delta pose (dx, dy, dz, droll, dpitch, dyaw). pub ee_delta: Option<[f32; 6]>, /// Gripper command (0 = open, 1 = closed). pub gripper: Option, } /// Transition tuple for training. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Transition { /// Current observation. pub observation: Observation, /// Action taken. pub action: Action, /// Reward received. pub reward: f32, /// Next observation. pub next_observation: Observation, /// Whether episode terminated. pub done: bool, /// Whether episode was truncated. pub truncated: bool, /// Additional info. pub info: TransitionInfo, } /// Additional transition information. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct TransitionInfo { /// Task success. pub success: Option, /// Distance to goal. pub distance_to_goal: Option, /// Contact forces. pub contact_forces: Option>, } // ============================================================================ // World Model Types // ============================================================================ /// World model type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum WorldModelType { /// RSSM (Recurrent State Space Model) - DreamerV3 style. RSSM, /// JEPA (Joint Embedding Predictive Architecture). JEPA, /// Transformer-based world model. Transformer, /// Latent Dynamics Model. LatentDynamics, /// Variational World Model. VWM, } /// World model configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WorldModelConfig { /// Model type. pub model_type: WorldModelType, /// Observation encoding dimension. pub obs_embed_dim: usize, /// Action encoding dimension. pub action_embed_dim: usize, /// Deterministic state dimension (RSSM). pub deter_dim: usize, /// Stochastic state dimension (RSSM). pub stoch_dim: usize, /// Number of stochastic classes (discrete). pub stoch_classes: usize, /// Hidden dimension for MLPs. pub hidden_dim: usize, /// Number of transformer layers (if applicable). pub num_layers: usize, /// Image encoder type. pub image_encoder: ImageEncoderType, /// Whether to predict rewards. pub predict_reward: bool, /// Whether to predict termination. pub predict_done: bool, /// Imagination horizon. pub imagination_horizon: usize, } impl Default for WorldModelConfig { fn default() -> Self { Self { 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, } } } /// Image encoder type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum ImageEncoderType { /// Convolutional neural network. CNN, /// Vision Transformer. ViT, /// ResNet-style. ResNet, /// Simple MLP (for vector observations). MLP, } // ============================================================================ // Policy Types // ============================================================================ /// Policy type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum PolicyType { /// Actor-Critic. ActorCritic, /// Model Predictive Control. MPC, /// Cross-Entropy Method. CEM, /// Behavioral Cloning. BehavioralCloning, /// Diffusion Policy. DiffusionPolicy, } /// Policy configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PolicyConfig { /// Policy type. pub policy_type: PolicyType, /// State dimension (from world model). pub state_dim: usize, /// Action dimension. pub action_dim: usize, /// Hidden dimensions for actor. pub actor_hidden: Vec, /// Hidden dimensions for critic. pub critic_hidden: Vec, /// Discount factor. pub gamma: f32, /// GAE lambda. pub gae_lambda: f32, /// Entropy coefficient. pub entropy_coef: f32, /// Use continuous actions. pub continuous_actions: bool, /// Action distribution type. pub action_dist: ActionDistribution, } impl Default for PolicyConfig { fn default() -> Self { Self { policy_type: PolicyType::ActorCritic, state_dim: 512, action_dim: 7, actor_hidden: vec![256, 256], critic_hidden: vec![256, 256], gamma: 0.99, gae_lambda: 0.95, entropy_coef: 0.001, continuous_actions: true, action_dist: ActionDistribution::Normal, } } } /// Action distribution type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum ActionDistribution { /// Normal (Gaussian) distribution. Normal, /// Truncated normal distribution. TruncatedNormal, /// Categorical distribution. Categorical, /// Mixture of Gaussians. MixtureOfGaussians, } // ============================================================================ // Training Types // ============================================================================ /// Training configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TrainingConfig { /// Batch size. pub batch_size: usize, /// Sequence length for world model. pub sequence_length: usize, /// Learning rate for world model. pub wm_learning_rate: f64, /// Learning rate for actor. pub actor_learning_rate: f64, /// Learning rate for critic. pub critic_learning_rate: f64, /// Number of training epochs. pub epochs: usize, /// Buffer size for replay. pub buffer_size: usize, /// Prefill steps. pub prefill_steps: usize, /// Train every N steps. pub train_every: usize, /// Train ratio (train steps per env step). pub train_ratio: f32, /// Use mixed precision. pub mixed_precision: bool, /// Gradient clipping. pub gradient_clip: f32, /// Random seed. pub seed: Option, } impl Default for TrainingConfig { fn default() -> Self { Self { batch_size: 64, sequence_length: 50, 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), } } } /// Training progress. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TrainingProgress { /// Current epoch. pub epoch: usize, /// Total epochs. pub total_epochs: usize, /// Current step. pub step: usize, /// Total steps. pub total_steps: usize, /// World model loss. pub wm_loss: f64, /// Image reconstruction loss. pub image_loss: Option, /// Reward prediction loss. pub reward_loss: Option, /// KL divergence loss. pub kl_loss: Option, /// Actor loss. pub actor_loss: Option, /// Critic loss. pub critic_loss: Option, /// Entropy. pub entropy: Option, /// Learning rate. pub learning_rate: f64, /// Environment return (evaluation). pub env_return: Option, /// Episode length (evaluation). pub episode_length: Option, /// Elapsed time in seconds. pub elapsed_seconds: f64, } // ============================================================================ // Environment Types // ============================================================================ /// Robot environment type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum EnvironmentType { /// Panda Franka robot. Panda, /// UR5 robot. UR5, /// Sawyer robot. Sawyer, /// Kuka robot. Kuka, /// Custom robot. Custom, } /// Task type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum TaskType { /// Reach a target position. Reach, /// Push an object. Push, /// Pick and place. PickPlace, /// Stack objects. Stack, /// Insert peg into hole. PegInsertion, /// Open a door. DoorOpen, /// General manipulation. Manipulation, } /// Environment configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EnvironmentConfig { /// Environment type. pub env_type: EnvironmentType, /// Task type. pub task_type: TaskType, /// Number of joints. pub num_joints: usize, /// Action dimension. pub action_dim: usize, /// Observation dimension (non-image). pub obs_dim: usize, /// Use camera observations. pub use_camera: bool, /// Camera image size. pub camera_size: (usize, usize), /// Maximum episode length. pub max_episode_length: usize, /// Action repeat. pub action_repeat: usize, /// Control frequency (Hz). pub control_freq: f32, /// Reward type. pub reward_type: RewardType, } impl Default for EnvironmentConfig { fn default() -> Self { Self { env_type: EnvironmentType::Panda, task_type: TaskType::Reach, num_joints: 7, action_dim: 7, obs_dim: 21, // 7 pos + 7 vel + 7 ee pose use_camera: true, camera_size: (64, 64), max_episode_length: 500, action_repeat: 2, control_freq: 20.0, reward_type: RewardType::Dense, } } } /// Reward type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum RewardType { /// Dense reward (continuous feedback). Dense, /// Sparse reward (only on success). Sparse, /// Shaped reward (curriculum). Shaped, } // ============================================================================ // Evaluation Types // ============================================================================ /// Evaluation request. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EvaluationRequest { /// Number of episodes. pub num_episodes: usize, /// Whether to render. pub render: bool, /// Use deterministic policy. pub deterministic: bool, /// Record video. pub record_video: bool, /// Environment config. pub env_config: EnvironmentConfig, } impl Default for EvaluationRequest { fn default() -> Self { Self { num_episodes: 10, render: false, deterministic: true, record_video: false, env_config: EnvironmentConfig::default(), } } } /// Evaluation result. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EvaluationResult { /// Mean return. pub mean_return: f64, /// Std return. pub std_return: f64, /// Mean episode length. pub mean_length: f64, /// Success rate. pub success_rate: f64, /// Individual episode returns. pub episode_returns: Vec, /// Individual episode lengths. pub episode_lengths: Vec, /// Total evaluation time. pub eval_time: f64, } // ============================================================================ // Imagination Types // ============================================================================ /// Imagined trajectory. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ImaginedTrajectory { /// States (world model latent states). pub states: Vec>, /// Actions taken. pub actions: Vec>, /// Predicted rewards. pub rewards: Vec, /// Predicted values. pub values: Vec, /// Predicted termination probabilities. pub dones: Vec, /// Decoded observations (optional). pub decoded_observations: Option>>, } /// Imagination request. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ImaginationRequest { /// Starting observation. pub initial_observation: Observation, /// Horizon length. pub horizon: usize, /// Number of trajectories. pub num_trajectories: usize, /// Use policy for actions. pub use_policy: bool, /// Decode observations. pub decode: bool, } // ============================================================================ // Sim-to-Real Types // ============================================================================ /// Domain randomization config. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DomainRandomization { /// Randomize dynamics. pub dynamics: bool, /// Randomize visuals. pub visuals: bool, /// Randomize lighting. pub lighting: bool, /// Mass range (multiplier). pub mass_range: (f32, f32), /// Friction range. pub friction_range: (f32, f32), /// Damping range. pub damping_range: (f32, f32), /// Action noise std. pub action_noise: f32, /// Observation noise std. pub observation_noise: f32, } impl Default for DomainRandomization { fn default() -> Self { Self { dynamics: true, visuals: true, lighting: true, mass_range: (0.8, 1.2), friction_range: (0.5, 1.5), damping_range: (0.8, 1.2), action_noise: 0.01, observation_noise: 0.01, } } } /// Sim-to-real transfer config. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SimToRealConfig { /// Domain randomization settings. pub domain_randomization: DomainRandomization, /// Use real robot calibration. pub use_calibration: bool, /// Action scaling factor. pub action_scale: f32, /// Safety limits. pub safety_limits: SafetyLimits, } /// Safety limits for real robot. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SafetyLimits { /// Maximum joint velocity (rad/s). pub max_joint_velocity: f32, /// Maximum joint acceleration (rad/s²). pub max_joint_acceleration: f32, /// Maximum end-effector velocity (m/s). pub max_ee_velocity: f32, /// Maximum force (N). pub max_force: f32, /// Workspace limits (min, max) for x, y, z. pub workspace: [[f32; 2]; 3], } impl Default for SafetyLimits { fn default() -> Self { Self { max_joint_velocity: 1.5, max_joint_acceleration: 5.0, max_ee_velocity: 0.5, max_force: 50.0, workspace: [ [-0.5, 0.5], // x [-0.5, 0.5], // y [0.0, 0.8], // z ], } } } // ============================================================================ // Sample Functions // ============================================================================ /// Create a sample observation. #[must_use] pub fn sample_observation() -> Observation { Observation { joint_positions: vec![0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785], joint_velocities: vec![0.0; 7], end_effector_pose: [0.3, 0.0, 0.5, 1.0, 0.0, 0.0, 0.0], camera_image: None, image_dims: None, timestamp: 0.0, } } /// Create a sample action. #[must_use] pub fn sample_action() -> Action { Action { joint_positions: Some(vec![0.1, -0.7, 0.0, -2.3, 0.0, 1.5, 0.8]), joint_velocities: None, joint_torques: None, ee_delta: None, gripper: Some(0.0), } } /// Create a sample transition. #[must_use] pub fn sample_transition() -> Transition { Transition { observation: sample_observation(), action: sample_action(), reward: 1.0, next_observation: sample_observation(), done: false, truncated: false, info: TransitionInfo { success: Some(false), distance_to_goal: Some(0.15), contact_forces: None, }, } } /// Create a sample world model config. #[must_use] pub fn sample_world_model_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 sample policy config. #[must_use] pub fn sample_policy_config() -> PolicyConfig { PolicyConfig { policy_type: PolicyType::ActorCritic, state_dim: 512 + 32 * 32, // deter + stoch * classes 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 a sample training config. #[must_use] pub fn sample_training_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: 50, buffer_size: 100_000, prefill_steps: 1000, train_every: 5, train_ratio: 256.0, mixed_precision: false, gradient_clip: 100.0, seed: Some(42), } } /// Create a sample environment config. #[must_use] pub fn sample_env_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, } } // ============================================================================ // Tests // ============================================================================ #[cfg(test)] mod tests { use super::*; #[test] fn test_observation() { let obs = sample_observation(); assert_eq!(obs.joint_positions.len(), 7); assert_eq!(obs.joint_velocities.len(), 7); } #[test] fn test_action() { let action = sample_action(); assert!(action.joint_positions.is_some()); assert!(action.gripper.is_some()); } #[test] fn test_transition() { let trans = sample_transition(); assert!(!trans.done); assert!(trans.info.distance_to_goal.is_some()); } #[test] fn test_world_model_config() { let config = sample_world_model_config(); assert_eq!(config.model_type, WorldModelType::RSSM); assert!(config.predict_reward); } #[test] fn test_policy_config() { let config = sample_policy_config(); assert_eq!(config.policy_type, PolicyType::ActorCritic); assert!(config.continuous_actions); } #[test] fn test_training_config() { let config = sample_training_config(); assert!(config.batch_size > 0); assert!(config.sequence_length > 0); } #[test] fn test_env_config() { let config = sample_env_config(); assert_eq!(config.env_type, EnvironmentType::Panda); assert_eq!(config.task_type, TaskType::Reach); } #[test] fn test_domain_randomization() { let dr = DomainRandomization::default(); assert!(dr.dynamics); assert!(dr.mass_range.0 < dr.mass_range.1); } #[test] fn test_safety_limits() { let limits = SafetyLimits::default(); assert!(limits.max_joint_velocity > 0.0); assert!(limits.max_force > 0.0); } #[test] fn test_serialization() { let config = sample_world_model_config(); let json = serde_json::to_string(&config).unwrap(); let parsed: WorldModelConfig = serde_json::from_str(&json).unwrap(); assert_eq!(parsed.model_type, config.model_type); } }