r2l-core 0.0.3

use crate::{
    HookResult, break_on_hook_result,
    buffers::{TrajectoryBatch, buffer::TrajectoryView},
    error::Error,
    models::Actor,
    return_on_hook_result,
    tensor::R2lTensor,
    utils::{actor_wrapper::ActorWrapper, buffer_wrapper::TrajectoryViewsWrapper},
};

/// Trainable on-policy component that owns an actor and learns from rollouts.
pub trait Agent {
    /// Tensor type shared with the sampler and rollout buffers.
    type Tensor: R2lTensor;

    /// Actor type used by samplers to collect new rollouts.
    type Actor: Actor<Tensor = Self::Tensor> + Clone;

    /// Returns an actor snapshot for rollout collection.
    fn actor(&self) -> Self::Actor;

    /// Learns from a batch of trajectory containers.
    ///
    /// # Errors
    ///
    /// Returns an error if the agent update fails.
    fn learn<B: TrajectoryBatch<Self::Tensor>>(&mut self, buffers: &[B]) -> Result<(), Error>;

    /// Sets the learning rate used by future updates.
    fn set_learning_rate(&mut self, learning_rate: f64);
}

/// Rollout collector used by an on-policy training loop.
pub trait Sampler {
    /// Tensor type stored in collected trajectories.
    type Tensor: R2lTensor;

    /// Resets all environments managed by the sampler.
    ///
    /// # Errors
    ///
    /// Returns an error if an environment cannot be reset.
    fn reset_all_envs(&mut self) -> Result<(), Error> {
        Ok(())
    }

    /// Collects rollout data using the provided actor.
    ///
    /// # Errors
    ///
    /// Returns an error if an environment operation fails during collection.
    fn collect_rollouts<A: Actor<Tensor = Self::Tensor> + Clone>(
        &mut self,
        actor: A,
    ) -> Result<(), Error>;

    /// Creates a view for the agents.
    fn trajectory_views(&mut self) -> impl AsRef<[TrajectoryView<'_, Self::Tensor>]>;
}

/// Coupled runtime unit that binds an agent and sampler together.
pub struct OnPolicyRuntime<A: Agent, S: Sampler> {
    /// Trainable agent.
    pub agent: A,
    /// Rollout collector.
    pub sampler: S,
}

impl<A: Agent, S: Sampler> OnPolicyRuntime<A, S> {
    /// Collects a fresh set of rollouts using the sampler-facing actor.
    ///
    /// # Errors
    ///
    /// Returns an error if the sampler cannot collect the rollouts.
    pub fn collect(&mut self) -> Result<(), Error> {
        let actor = self.agent.actor();
        let actor = ActorWrapper::new(actor);
        self.sampler.collect_rollouts(actor)
    }

    /// Returns the last collected trajectory containers from the sampler.
    pub fn trajectory_containers(&mut self) -> impl AsRef<[TrajectoryView<'_, S::Tensor>]> {
        self.sampler.trajectory_views()
    }

    /// Adapts the sampler buffers and runs an agent update.
    ///
    /// # Errors
    ///
    /// Returns an error if the agent cannot learn from the collected trajectories.
    pub fn learn(&mut self) -> Result<(), Error> {
        let views = self.sampler.trajectory_views();
        let buffers = views
            .as_ref()
            .iter()
            .map(TrajectoryViewsWrapper::from_view)
            .collect::<Result<Vec<_>, _>>()?;
        self.agent.learn(&buffers)
    }

    /// Returns the agent-facing actor snapshot.
    pub fn actor(&self) -> A::Actor {
        self.agent.actor()
    }
}

/// Lifecycle hooks that control an [`OnPolicyAlgorithm`] training loop.
pub trait OnPolicyAlgorithmHooks {
    /// Agent type controlled by the training loop.
    type A: Agent;
    /// Sampler type controlled by the training loop.
    type S: Sampler;

    /// Called once before rollout/training starts.
    fn init_hook(&mut self, runtime: &mut OnPolicyRuntime<Self::A, Self::S>) -> HookResult;

    /// Called after rollouts are collected and before agent learning.
    fn post_rollout_hook(&mut self, runtime: &mut OnPolicyRuntime<Self::A, Self::S>) -> HookResult;

    /// Called after the agent has learned from the latest rollouts.
    fn post_training_hook(&mut self, runtime: &mut OnPolicyRuntime<Self::A, Self::S>)
    -> HookResult;

    /// Called once when the loop exits.
    ///
    /// # Errors
    ///
    /// Returns an error if end-of-training finalization fails.
    fn finish_training_hook(
        &mut self,
        runtime: &mut OnPolicyRuntime<Self::A, Self::S>,
    ) -> Result<(), Error>;
}

/// Generic on-policy training loop combining a runtime with lifecycle hooks.
pub struct OnPolicyAlgorithm<A: Agent, S: Sampler, H: OnPolicyAlgorithmHooks<A = A, S = S>> {
    /// Coupled training runtime.
    pub runtime: OnPolicyRuntime<A, S>,
    /// Lifecycle hooks.
    pub hooks: H,
}

impl<A: Agent, S: Sampler, H: OnPolicyAlgorithmHooks<A = A, S = S>> OnPolicyAlgorithm<A, S, H> {
    fn training_loop(&mut self) -> Result<(), Error> {
        return_on_hook_result!(self.hooks.init_hook(&mut self.runtime));
        loop {
            self.runtime.collect()?;
            break_on_hook_result!(self.hooks.post_rollout_hook(&mut self.runtime));

            self.runtime.learn()?;
            break_on_hook_result!(self.hooks.post_training_hook(&mut self.runtime));
        }
        Ok(())
    }
}

impl<A: Agent, S: Sampler, H: OnPolicyAlgorithmHooks<A = A, S = S>> OnPolicyAlgorithm<A, S, H> {
    /// Creates an on-policy algorithm from its runtime and lifecycle hooks.
    pub fn new(runtime: OnPolicyRuntime<A, S>, hooks: H) -> Self {
        Self { runtime, hooks }
    }

    /// Runs training until a hook requests termination.
    ///
    /// # Errors
    ///
    /// Returns an error if learning fails or a hook reports a deferred failure
    /// during end-of-training finalization.
    pub fn train(&mut self) -> Result<(), Error> {
        let training_result = self.training_loop();
        let finalization_result = self.hooks.finish_training_hook(&mut self.runtime);
        training_result.and(finalization_result)
    }
}