quantrs2_ml/

reinforcement.rs

//! Quantum Reinforcement Learning (QRL) agents and environments.
//!
//! Defines the [`Environment`] trait plus [`QuantumRLAgent`] which uses a
//! quantum neural network as its policy / value function approximator,
//! trained via policy gradient or Q-learning objectives.

use crate::error::{MLError, Result};
use crate::qnn::QuantumNeuralNetwork;
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::prelude::*;
use std::collections::HashMap;

/// Environment for reinforcement learning
pub trait Environment {
    /// Gets the current state
    fn state(&self) -> Array1<f64>;

    /// Gets the number of available actions
    fn num_actions(&self) -> usize;

    /// Takes an action and returns the reward and next state
    fn step(&mut self, action: usize) -> Result<(Array1<f64>, f64, bool)>;

    /// Resets the environment
    fn reset(&mut self) -> Array1<f64>;
}

/// Agent for reinforcement learning
pub trait QuantumAgent {
    /// Gets an action for a given state
    fn get_action(&self, state: &Array1<f64>) -> Result<usize>;

    /// Updates the agent based on a reward
    fn update(
        &mut self,
        state: &Array1<f64>,
        action: usize,
        reward: f64,
        next_state: &Array1<f64>,
        done: bool,
    ) -> Result<()>;

    /// Trains the agent on an environment
    fn train(&mut self, env: &mut dyn Environment, episodes: usize) -> Result<f64>;

    /// Evaluates the agent on an environment
    fn evaluate(&self, env: &mut dyn Environment, episodes: usize) -> Result<f64>;
}

/// Reinforcement learning algorithm type
#[derive(Debug, Clone, Copy)]
pub enum ReinforcementLearningType {
    /// Q-learning
    QLearning,

    /// SARSA
    SARSA,

    /// Deep Q-Network
    DQN,

    /// Policy Gradient
    PolicyGradient,

    /// Quantum Approximate Optimization Algorithm
    QAOA,
}

/// Reinforcement learning with quantum circuit
#[derive(Debug, Clone)]
pub struct ReinforcementLearning {
    /// Type of reinforcement learning algorithm
    rl_type: ReinforcementLearningType,

    /// Quantum neural network
    qnn: QuantumNeuralNetwork,

    /// Learning rate
    learning_rate: f64,

    /// Discount factor
    discount_factor: f64,

    /// Exploration rate
    exploration_rate: f64,

    /// Number of state dimensions
    state_dim: usize,

    /// Number of actions
    action_dim: usize,
}

impl ReinforcementLearning {
    /// Creates a new quantum reinforcement learning agent
    ///
    /// # Errors
    /// Returns an error if the quantum neural network cannot be created
    pub fn new() -> Result<Self> {
        // This is a placeholder implementation
        // In a real system, this would create a proper QNN

        let layers = vec![
            crate::qnn::QNNLayerType::EncodingLayer { num_features: 4 },
            crate::qnn::QNNLayerType::VariationalLayer { num_params: 16 },
            crate::qnn::QNNLayerType::EntanglementLayer {
                connectivity: "full".to_string(),
            },
            crate::qnn::QNNLayerType::VariationalLayer { num_params: 16 },
            crate::qnn::QNNLayerType::MeasurementLayer {
                measurement_basis: "computational".to_string(),
            },
        ];

        let qnn = QuantumNeuralNetwork::new(
            layers, 8, // 8 qubits
            4, // 4 input features
            2, // 2 output actions
        )?;

        Ok(ReinforcementLearning {
            rl_type: ReinforcementLearningType::QLearning,
            qnn,
            learning_rate: 0.01,
            discount_factor: 0.95,
            exploration_rate: 0.1,
            state_dim: 4,
            action_dim: 2,
        })
    }

    /// Sets the reinforcement learning algorithm type
    pub fn with_algorithm(mut self, rl_type: ReinforcementLearningType) -> Self {
        self.rl_type = rl_type;
        self
    }

    /// Sets the state dimension
    pub fn with_state_dimension(mut self, state_dim: usize) -> Self {
        self.state_dim = state_dim;
        self
    }

    /// Sets the action dimension
    pub fn with_action_dimension(mut self, action_dim: usize) -> Self {
        self.action_dim = action_dim;
        self
    }

    /// Sets the learning rate
    pub fn with_learning_rate(mut self, learning_rate: f64) -> Self {
        self.learning_rate = learning_rate;
        self
    }

    /// Sets the discount factor
    pub fn with_discount_factor(mut self, discount_factor: f64) -> Self {
        self.discount_factor = discount_factor;
        self
    }

    /// Sets the exploration rate
    pub fn with_exploration_rate(mut self, exploration_rate: f64) -> Self {
        self.exploration_rate = exploration_rate;
        self
    }

    /// Gets the Q-values for a state by evaluating the quantum neural network.
    ///
    /// The state is encoded into the QNN's parameterised circuit, simulated on
    /// the state-vector backend, and each action's Q-value is read out as a
    /// Pauli expectation value (one output per action).
    fn get_q_values(&self, state: &Array1<f64>) -> Result<Array1<f64>> {
        self.qnn.forward(state)
    }
}

impl QuantumAgent for ReinforcementLearning {
    fn get_action(&self, state: &Array1<f64>) -> Result<usize> {
        // Epsilon-greedy action selection
        if thread_rng().random::<f64>() < self.exploration_rate {
            // Explore: random action
            Ok(fastrand::usize(0..self.action_dim))
        } else {
            // Exploit: best action according to the QNN Q-values.
            let q_values = self.get_q_values(state)?;
            if q_values.is_empty() {
                return Err(MLError::MLOperationError(
                    "QNN produced no Q-values".to_string(),
                ));
            }
            let mut best_action = 0;
            for i in 1..q_values.len() {
                if q_values[i] > q_values[best_action] {
                    best_action = i;
                }
            }

            Ok(best_action)
        }
    }

    fn update(
        &mut self,
        state: &Array1<f64>,
        action: usize,
        reward: f64,
        next_state: &Array1<f64>,
        done: bool,
    ) -> Result<()> {
        // Semi-gradient Q-learning update on the quantum neural network.
        let q_values = self.get_q_values(state)?;
        if action >= q_values.len() {
            return Err(MLError::InvalidParameter(format!(
                "action {action} out of range for {} Q-values",
                q_values.len()
            )));
        }
        let q_sa = q_values[action];

        // Bellman target: r + γ · max_a' Q(s', a')  (terminal transitions use r).
        let target = if done {
            reward
        } else {
            let next_q = self.get_q_values(next_state)?;
            let max_next = next_q.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            reward + self.discount_factor * max_next
        };

        // Loss L = (Q(s,a) - target)²  ⇒  dL/dθ = 2·(Q(s,a) - target)·∂Q(s,a)/∂θ,
        // where ∂Q(s,a)/∂θ is obtained by the parameter-shift rule.
        let td_error = q_sa - target;
        let gradient = self.qnn.output_component_gradient(state, action)?;
        for j in 0..self.qnn.parameters.len() {
            self.qnn.parameters[j] -= self.learning_rate * 2.0 * td_error * gradient[j];
        }

        Ok(())
    }

    fn train(&mut self, env: &mut dyn Environment, episodes: usize) -> Result<f64> {
        let mut total_reward = 0.0;

        for _ in 0..episodes {
            let mut state = env.reset();
            let mut episode_reward = 0.0;
            let mut done = false;

            while !done {
                let action = self.get_action(&state)?;
                let (next_state, reward, is_done) = env.step(action)?;

                self.update(&state, action, reward, &next_state, is_done)?;

                state = next_state;
                episode_reward += reward;
                done = is_done;
            }

            total_reward += episode_reward;
        }

        Ok(total_reward / episodes as f64)
    }

    fn evaluate(&self, env: &mut dyn Environment, episodes: usize) -> Result<f64> {
        let mut total_reward = 0.0;

        for _ in 0..episodes {
            let mut state = env.reset();
            let mut episode_reward = 0.0;
            let mut done = false;

            while !done {
                let action = self.get_action(&state)?;
                let (next_state, reward, is_done) = env.step(action)?;

                state = next_state;
                episode_reward += reward;
                done = is_done;
            }

            total_reward += episode_reward;
        }

        Ok(total_reward / episodes as f64)
    }
}

/// GridWorld environment for testing reinforcement learning
pub struct GridWorldEnvironment {
    /// Width of the grid
    width: usize,

    /// Height of the grid
    height: usize,

    /// Current position (x, y)
    position: (usize, usize),

    /// Goal position (x, y)
    goal: (usize, usize),

    /// Obstacle positions (x, y)
    obstacles: Vec<(usize, usize)>,
}

impl GridWorldEnvironment {
    /// Creates a new GridWorld environment
    pub fn new(width: usize, height: usize) -> Self {
        GridWorldEnvironment {
            width,
            height,
            position: (0, 0),
            goal: (width - 1, height - 1),
            obstacles: Vec::new(),
        }
    }

    /// Sets the goal position
    pub fn with_goal(mut self, x: usize, y: usize) -> Self {
        self.goal = (x.min(self.width - 1), y.min(self.height - 1));
        self
    }

    /// Sets the obstacles
    pub fn with_obstacles(mut self, obstacles: Vec<(usize, usize)>) -> Self {
        self.obstacles = obstacles;
        self
    }

    /// Checks if a position is an obstacle
    pub fn is_obstacle(&self, x: usize, y: usize) -> bool {
        self.obstacles.contains(&(x, y))
    }

    /// Checks if a position is the goal
    pub fn is_goal(&self, x: usize, y: usize) -> bool {
        (x, y) == self.goal
    }
}

impl Environment for GridWorldEnvironment {
    fn state(&self) -> Array1<f64> {
        let mut state = Array1::zeros(4);

        // Normalize position
        state[0] = self.position.0 as f64 / self.width as f64;
        state[1] = self.position.1 as f64 / self.height as f64;

        // Normalize goal
        state[2] = self.goal.0 as f64 / self.width as f64;
        state[3] = self.goal.1 as f64 / self.height as f64;

        state
    }

    fn num_actions(&self) -> usize {
        4 // Up, Right, Down, Left
    }

    fn step(&mut self, action: usize) -> Result<(Array1<f64>, f64, bool)> {
        // Calculate new position
        let (x, y) = self.position;
        let (new_x, new_y) = match action {
            0 => (x, y.saturating_sub(1)), // Up
            1 => (x + 1, y),               // Right
            2 => (x, y + 1),               // Down
            3 => (x.saturating_sub(1), y), // Left
            _ => {
                return Err(MLError::InvalidParameter(format!(
                    "Invalid action: {}",
                    action
                )))
            }
        };

        // Check if new position is valid
        let new_x = new_x.min(self.width - 1);
        let new_y = new_y.min(self.height - 1);

        // Check if new position is an obstacle
        if self.obstacles.contains(&(new_x, new_y)) {
            // Stay in the same position
            let reward = -1.0;
            let done = false;
            return Ok((self.state(), reward, done));
        }

        // Update position
        self.position = (new_x, new_y);

        // Calculate reward
        let reward = if (new_x, new_y) == self.goal {
            10.0 // Goal reached
        } else {
            -0.1 // Step penalty
        };

        // Check if done
        let done = (new_x, new_y) == self.goal;

        Ok((self.state(), reward, done))
    }

    fn reset(&mut self) -> Array1<f64> {
        self.position = (0, 0);
        self.state()
    }
}