burn-train 0.22.0-pre.1

Training crate for the Burn 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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::marker::PhantomData;

use burn_rl::{
    Batchable, Environment, EnvironmentInit, Policy, PolicyLearner, PolicyState, ToAction,
    ToObservation,
};

use crate::checkpoint::Checkpoint;
use crate::{AgentEvaluationEvent, AsyncProcessorTraining, ItemLazy, RLEvent};

/// All components used by the reinforcement learning paradigm, grouped in one trait.
pub trait RLComponentsTypes {
    /// The learning environment.
    type Env: Environment<State = Self::State, Action = Self::Action> + 'static;
    /// Specifies how to initialize the environment.
    type EnvInit: EnvironmentInit<Self::Env> + Send + 'static;
    /// The type of the environment state.
    type State: ToObservation<<Self::Policy as Policy>::Observation> + Clone + Send + 'static;
    /// The type of the environment action.
    type Action: From<<Self::Policy as Policy>::Action>
        + ToAction<<Self::Policy as Policy>::Action>
        + Clone
        + Send
        + 'static;

    /// The policy used to take actions in the environment.
    type Policy: Policy<
            Observation = Self::PolicyObs,
            ActionDistribution = Self::PolicyAD,
            Action = Self::PolicyAction,
            ActionContext = Self::ActionContext,
            PolicyState = Self::PolicyState,
        > + Send
        + 'static;
    /// The policy's observation type.
    type PolicyObs: Clone + Send + Batchable + 'static;
    /// The policy's action distribution type.
    type PolicyAD: Clone + Send + Batchable;
    /// The policy's action type.
    type PolicyAction: Clone + Send + Batchable;
    /// Additional data as context for an agent's action.
    type ActionContext: ItemLazy + Clone + Send + 'static;
    /// The state of the parameterized policy.
    type PolicyState: Clone + Send + PolicyState<Record: Checkpoint> + 'static;

    /// The learning agent.
    type LearningAgent: PolicyLearner<
            TrainContext = Self::TrainingOutput,
            InnerPolicy = Self::Policy,
            Record: Checkpoint,
        > + Send
        + 'static;
    /// The output data of a training step.
    type TrainingOutput: ItemLazy + Clone + Send;
}

/// Concrete type that implements the [RLComponentsTypes](RLComponentsTypes) trait.
pub struct RLComponentsMarker<E, EI, A> {
    _env: PhantomData<E>,
    _env_init: PhantomData<EI>,
    _agent: PhantomData<A>,
}

impl<E, EI, A> RLComponentsTypes for RLComponentsMarker<E, EI, A>
where
    E: Environment + 'static,
    EI: EnvironmentInit<E> + Send + 'static,
    A: PolicyLearner + Send + 'static,
    <A as PolicyLearner>::Record: Checkpoint,
    A::TrainContext: ItemLazy + Clone + Send,
    A::InnerPolicy: Policy + Send,
    <A::InnerPolicy as Policy>::Observation: Batchable + Clone + Send,
    <A::InnerPolicy as Policy>::ActionDistribution: Batchable + Clone + Send,
    <A::InnerPolicy as Policy>::Action: Batchable + Clone + Send,
    <A::InnerPolicy as Policy>::ActionContext: ItemLazy + Clone + Send + 'static,
    <A::InnerPolicy as Policy>::PolicyState: Clone + Send,
    <<A::InnerPolicy as Policy>::PolicyState as PolicyState>::Record: Checkpoint,
    E::State: ToObservation<<A::InnerPolicy as Policy>::Observation> + Clone + Send + 'static,
    E::Action: From<<A::InnerPolicy as Policy>::Action>
        + ToAction<<A::InnerPolicy as Policy>::Action>
        + Clone
        + Send
        + 'static,
{
    type Env = E;
    type EnvInit = EI;
    type LearningAgent = A;
    type Policy = A::InnerPolicy;
    type PolicyObs = <A::InnerPolicy as Policy>::Observation;
    type PolicyAD = <A::InnerPolicy as Policy>::ActionDistribution;
    type PolicyAction = <A::InnerPolicy as Policy>::Action;
    type ActionContext = <A::InnerPolicy as Policy>::ActionContext;
    type PolicyState = <A::InnerPolicy as Policy>::PolicyState;
    type TrainingOutput = A::TrainContext;
    type State = E::State;
    type Action = E::Action;
}

pub(crate) type RlPolicy<RLC> =
    <<RLC as RLComponentsTypes>::LearningAgent as PolicyLearner>::InnerPolicy;
/// The event processor type for reinforcement learning.
pub type RLEventProcessorType<RLC> = AsyncProcessorTraining<
    RLEvent<<RLC as RLComponentsTypes>::TrainingOutput, <RLC as RLComponentsTypes>::ActionContext>,
    AgentEvaluationEvent<<RLC as RLComponentsTypes>::ActionContext>,
>;
/// The record of the policy.
pub type RLPolicyRecord<RLC> =
    <<<RLC as RLComponentsTypes>::Policy as Policy>::PolicyState as PolicyState>::Record;
/// The record of the learning agent.
pub type RLAgentRecord<RLC> = <<RLC as RLComponentsTypes>::LearningAgent as PolicyLearner>::Record;