burn-rl 0.22.0-pre.2

use derive_new::new;

use burn_core::tensor::Device;

use crate::TransitionBatch;

/// An action along with additional context about the decision.
#[derive(Clone, new)]
pub struct ActionContext<A, C> {
    /// The context.
    pub context: C,
    /// The action.
    pub action: A,
}

/// The state of a policy.
pub trait PolicyState {
    /// The type of the record.
    type Record;

    /// Convert the state to a record.
    fn into_record(self) -> Self::Record;
    /// Load the state from a record.
    fn load_record(&self, record: Self::Record) -> Self;
}

/// Defines how an environment's state is converted to a policy's observation.
pub trait ToObservation<O> {
    /// Convert an environment's state to a policy's observation, moving it to the given device if needed.
    fn to_observation(&self, device: &Device) -> O;
}

/// Defines how an environment's action is converted to a policy's action.
pub trait ToAction<A> {
    /// Convert an environment's action to a policy's action, moving it to the given device if needed.
    fn to_action(&self, device: &Device) -> A;
}

/// Trait for a RL policy.
pub trait Policy: Clone {
    /// The observation given as input to the policy.
    type Observation;
    /// The action distribution parameters defining how the action will be sampled.
    type ActionDistribution;
    /// The action.
    type Action;

    /// Additional context on the policy's decision.
    type ActionContext;
    /// The current parameterization of the policy.
    type PolicyState: PolicyState;

    /// Produces the action distribution from a batch of observations.
    fn forward(&mut self, obs: Self::Observation) -> Self::ActionDistribution;
    /// Gives the action from a batch of observations.
    fn action(
        &mut self,
        obs: Self::Observation,
        deterministic: bool,
    ) -> (Self::Action, Vec<Self::ActionContext>);

    /// Update the policy's parameters.
    fn update(&mut self, update: Self::PolicyState);
    /// Returns the current parameterization.
    fn state(&self) -> Self::PolicyState;

    /// Loads the policy on the given device.
    fn to_device(self, device: &Device) -> Self;
    /// Loads the policy parameters from a record.
    fn load_record(self, record: <Self::PolicyState as PolicyState>::Record) -> Self;
}

/// Trait for a type that can be batched and unbatched (split).
pub trait Batchable: Sized {
    /// Create a batch from a list of items.
    fn batch(value: Vec<Self>) -> Self;
    /// Create a list from batched items.
    fn unbatch(self) -> Vec<Self>;
}

/// A training output.
pub struct RLTrainOutput<TO, P> {
    /// The policy.
    pub policy: P,
    /// The item.
    pub item: TO,
}

/// Batched transitions for a PolicyLearner.
pub type LearnerTransitionBatch<P> =
    TransitionBatch<<P as Policy>::Observation, <P as Policy>::Action>;

/// Learner for a policy.
pub trait PolicyLearner
where
    <Self::InnerPolicy as Policy>::Observation: Clone + Batchable,
    <Self::InnerPolicy as Policy>::ActionDistribution: Clone + Batchable,
    <Self::InnerPolicy as Policy>::Action: Clone + Batchable,
{
    /// Additional context of a training step.
    type TrainContext;
    /// The policy to train.
    type InnerPolicy: Policy;
    /// The record of the learner.
    type Record;

    /// Execute a training step on the policy.
    fn train(
        &mut self,
        input: LearnerTransitionBatch<Self::InnerPolicy>,
    ) -> RLTrainOutput<Self::TrainContext, <Self::InnerPolicy as Policy>::PolicyState>;
    /// Returns the learner's current policy for validation.
    fn policy(&self) -> Self::InnerPolicy;
    /// Update the learner's policy.
    fn update_policy(&mut self, update: Self::InnerPolicy);
    /// Convert the learner's state into a record.
    fn record(&self) -> Self::Record;
    /// Load the learner's state from a record.
    fn load_record(self, record: Self::Record) -> Self;
    /// Returns the device used for training.
    fn device(&self) -> Device;
}