relayrl_algorithms 0.4.1

Independent / Multi-Agent Proximal Policy Optimization (PPO) Algorithms for the RelayRL framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

use burn_tensor::backend::Backend;
use burn_tensor::{BasicOps, Tensor, TensorKind};
use relayrl_types::data::tensor::DType;
use relayrl_types::prelude::tensor::relayrl::BackendMatcher;

use super::types::{ArchLayer, LayerSpecs};

/// Combined supertrait for all algorithm-compatible networks: spec, forward pass, and weight export.
///
/// Implement this trait to plug a custom architecture into `PPOTrainerSpec`.
pub trait NeuralNetwork<B, KindIn, KindOut>:
    NeuralNetworkSpec<B, KindIn, KindOut> + NeuralNetworkForward<B, KindIn, KindOut> + WeightProvider
where
    B: Backend + BackendMatcher<Backend = B>,
    KindIn: TensorKind<B> + BasicOps<B>,
    KindOut: TensorKind<B> + BasicOps<B>,
{
    fn default(
        input_dim: usize,
        input_dtype: DType,
        output_dim: usize,
        output_dtype: DType,
        device: &B::Device,
    ) -> Self;
}

/// Exposes the input/output dimensions and dtypes of a network for validation and config.
pub trait NeuralNetworkSpec<
    B: Backend + BackendMatcher<Backend = B>,
    KindIn: TensorKind<B> + BasicOps<B>,
    KindOut: TensorKind<B> + BasicOps<B>,
>
{
    fn input_dim(&self) -> &usize;
    fn input_dtype(&self) -> &DType;
    fn output_dim(&self) -> &usize;
    fn output_dtype(&self) -> &DType;
}

/// Generic forward pass used by actors and algorithm kernels.
pub trait NeuralNetworkForward<
    B: Backend + BackendMatcher<Backend = B>,
    KindIn: TensorKind<B> + BasicOps<B>,
    KindOut: TensorKind<B> + BasicOps<B>,
>
{
    fn forward<const IN_D: usize, const OUT_D: usize>(
        &self,
        input: Tensor<B, IN_D, KindIn>,
    ) -> Tensor<B, OUT_D, KindOut>;
}

/// Trait for extracting per-layer weight specs from a network.
pub trait WeightProvider {
    fn get_layer_specs(&self) -> LayerSpecs;

    /// Return an architecture-aware layer sequence for conv / mixed networks.
    ///
    /// The default implementation returns `None`, preserving backward
    /// compatibility for all MLP implementations.  Override this for networks
    /// that cannot be represented as flat `LayerSpecs` (e.g. `ConvNetPolicy`).
    fn get_arch_spec(&self) -> Option<Vec<ArchLayer>> {
        None
    }
}