r2l_core/
models.rs

1use std::{fmt, str::FromStr};
2
3use serde::{Deserialize, Serialize};
4
5use crate::{error::Result, tensor::R2lTensor};
6
7/// Activation function used between hidden layers in feed-forward networks.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
9pub enum ActivationFunction {
10    /// Exponential linear unit activation with the backend default alpha.
11    Elu,
12    /// Gaussian error linear unit activation.
13    Gelu,
14    /// Gaussian error linear unit activation using the backend tanh approximation.
15    GeluApproximate,
16    /// Hard sigmoid activation with backend default parameters.
17    HardSigmoid,
18    /// Hard swish activation.
19    HardSwish,
20    /// Leaky rectified linear unit activation with the backend default slope.
21    LeakyRelu,
22    /// Rectified linear unit activation.
23    Relu,
24    /// Sigmoid activation.
25    Sigmoid,
26    /// Hyperbolic tangent activation.
27    #[default]
28    Tanh,
29}
30
31impl fmt::Display for ActivationFunction {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        let name = match self {
34            Self::Elu => "elu",
35            Self::Gelu => "gelu",
36            Self::GeluApproximate => "gelu_approximate",
37            Self::HardSigmoid => "hard_sigmoid",
38            Self::HardSwish => "hard_swish",
39            Self::LeakyRelu => "leaky_relu",
40            Self::Relu => "relu",
41            Self::Sigmoid => "sigmoid",
42            Self::Tanh => "tanh",
43        };
44        f.write_str(name)
45    }
46}
47
48impl FromStr for ActivationFunction {
49    type Err = String;
50
51    fn from_str(name: &str) -> std::result::Result<Self, Self::Err> {
52        match name {
53            "elu" => Ok(Self::Elu),
54            "gelu" => Ok(Self::Gelu),
55            "gelu_approximate" => Ok(Self::GeluApproximate),
56            "hard_sigmoid" => Ok(Self::HardSigmoid),
57            "hard_swish" => Ok(Self::HardSwish),
58            "leaky_relu" => Ok(Self::LeakyRelu),
59            "relu" => Ok(Self::Relu),
60            "sigmoid" => Ok(Self::Sigmoid),
61            "tanh" => Ok(Self::Tanh),
62            _ => Err(format!("unknown activation function: {name}")),
63        }
64    }
65}
66
67/// A policy-like object that can choose an action for one observation.
68///
69/// Actors are the inference-time surface used by samplers. They must be
70/// sendable so rollout collection can move them into worker threads.
71pub trait Actor: Send + 'static {
72    /// Tensor type accepted as observations and returned as actions.
73    type Tensor: R2lTensor;
74
75    /// Selects an action for a single observation.
76    ///
77    /// # Errors
78    ///
79    /// Returns an error if action inference fails.
80    fn action(&self, observation: Self::Tensor) -> Result<Self::Tensor>;
81
82    /// Selects the modal action for a single observation without sampling.
83    ///
84    /// # Errors
85    ///
86    /// Returns an error if action inference fails.
87    fn mode_action(&self, observation: Self::Tensor) -> Result<Self::Tensor>;
88}
89
90/// A policy that can be serialized as a safetensors artifact.
91pub trait ToSafetensors {
92    /// Serializes this policy as safetensors bytes.
93    ///
94    /// # Errors
95    ///
96    /// Returns an error if the policy parameters cannot be serialized.
97    fn to_safetensors(&self) -> Result<Vec<u8>>;
98}
99
100/// Trainable action distribution interface used by on-policy algorithms.
101///
102/// A `Policy` extends [`Actor`] with the quantities needed to compute policy
103/// gradient losses and entropy bonuses over a batch.
104pub trait Policy: Actor {
105    /// Computes log probabilities for batched observation/action pairs.
106    ///
107    /// # Errors
108    ///
109    /// Returns an error if the policy cannot evaluate the batch.
110    fn log_probs(
111        &self,
112        observations: &[Self::Tensor],
113        actions: &[Self::Tensor],
114    ) -> Result<Self::Tensor>;
115
116    /// Returns a representative action standard deviation when available.
117    ///
118    /// # Errors
119    ///
120    /// Returns an error if the standard deviation cannot be computed.
121    /// Returns `Ok(None)` when the policy has no meaningful scalar standard deviation.
122    fn std(&self) -> Result<Option<f32>>;
123
124    /// Computes the policy entropy for a batch of states.
125    ///
126    /// # Errors
127    ///
128    /// Returns an error if the entropy cannot be computed.
129    fn entropy(&self, states: &[Self::Tensor]) -> Result<Self::Tensor>;
130}
131
132/// Component that learns from backend-specific loss values.
133pub trait Learner {
134    /// Loss bundle consumed by this module.
135    type Losses;
136
137    /// Applies one optimization update from precomputed losses.
138    ///
139    /// # Errors
140    ///
141    /// Returns an error if the optimizer update fails.
142    fn update(&mut self, losses: Self::Losses) -> Result<()>;
143}
144
145/// Batched value-function interface.
146pub trait ValueFunction {
147    /// Tensor type used for observations and returned values.
148    type Tensor: Clone;
149
150    /// Estimates values for a batch of observations.
151    ///
152    /// # Errors
153    ///
154    /// Returns an error if value inference fails.
155    fn values(&self, observations: &[Self::Tensor]) -> Result<Self::Tensor>;
156}
r2l_core/models.rs

r2l_core/
models.rs
