rustytorch/demos/embodied-shared/src/lib.rs

//! 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<f32>,
    /// Joint velocities (rad/s).
    pub joint_velocities: Vec<f32>,
    /// End-effector pose (x, y, z, qw, qx, qy, qz).
    pub end_effector_pose: [f32; 7],
    /// Camera image (flattened RGB).
    pub camera_image: Option<Vec<u8>>,
    /// 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<Vec<f32>>,
    /// Joint velocity targets (rad/s).
    pub joint_velocities: Option<Vec<f32>>,
    /// Joint torques (Nm).
    pub joint_torques: Option<Vec<f32>>,
    /// 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<f32>,
}

/// 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<bool>,
    /// Distance to goal.
    pub distance_to_goal: Option<f32>,
    /// Contact forces.
    pub contact_forces: Option<Vec<f32>>,
}

// ============================================================================
// 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<usize>,
    /// Hidden dimensions for critic.
    pub critic_hidden: Vec<usize>,
    /// 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<u64>,
}

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<f64>,
    /// Reward prediction loss.
    pub reward_loss: Option<f64>,
    /// KL divergence loss.
    pub kl_loss: Option<f64>,
    /// Actor loss.
    pub actor_loss: Option<f64>,
    /// Critic loss.
    pub critic_loss: Option<f64>,
    /// Entropy.
    pub entropy: Option<f64>,
    /// Learning rate.
    pub learning_rate: f64,
    /// Environment return (evaluation).
    pub env_return: Option<f64>,
    /// Episode length (evaluation).
    pub episode_length: Option<f64>,
    /// 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<f64>,
    /// Individual episode lengths.
    pub episode_lengths: Vec<usize>,
    /// 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<Vec<f64>>,
    /// Actions taken.
    pub actions: Vec<Vec<f64>>,
    /// Predicted rewards.
    pub rewards: Vec<f64>,
    /// Predicted values.
    pub values: Vec<f64>,
    /// Predicted termination probabilities.
    pub dones: Vec<f64>,
    /// Decoded observations (optional).
    pub decoded_observations: Option<Vec<Vec<u8>>>,
}

/// 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);
    }
}