burn_train/learner/rl/
components.rs

1use std::marker::PhantomData;
2
3use burn_rl::{
4    Batchable, Environment, EnvironmentInit, Policy, PolicyLearner, PolicyState, ToAction,
5    ToObservation,
6};
7
8use crate::checkpoint::Checkpoint;
9use crate::{AgentEvaluationEvent, AsyncProcessorTraining, ItemLazy, RLEvent};
10
11/// All components used by the reinforcement learning paradigm, grouped in one trait.
12pub trait RLComponentsTypes {
13    /// The learning environment.
14    type Env: Environment<State = Self::State, Action = Self::Action> + 'static;
15    /// Specifies how to initialize the environment.
16    type EnvInit: EnvironmentInit<Self::Env> + Send + 'static;
17    /// The type of the environment state.
18    type State: ToObservation<<Self::Policy as Policy>::Observation> + Clone + Send + 'static;
19    /// The type of the environment action.
20    type Action: From<<Self::Policy as Policy>::Action>
21        + ToAction<<Self::Policy as Policy>::Action>
22        + Clone
23        + Send
24        + 'static;
25
26    /// The policy used to take actions in the environment.
27    type Policy: Policy<
28            Observation = Self::PolicyObs,
29            ActionDistribution = Self::PolicyAD,
30            Action = Self::PolicyAction,
31            ActionContext = Self::ActionContext,
32            PolicyState = Self::PolicyState,
33        > + Send
34        + 'static;
35    /// The policy's observation type.
36    type PolicyObs: Clone + Send + Batchable + 'static;
37    /// The policy's action distribution type.
38    type PolicyAD: Clone + Send + Batchable;
39    /// The policy's action type.
40    type PolicyAction: Clone + Send + Batchable;
41    /// Additional data as context for an agent's action.
42    type ActionContext: ItemLazy + Clone + Send + 'static;
43    /// The state of the parameterized policy.
44    type PolicyState: Clone + Send + PolicyState<Record: Checkpoint> + 'static;
45
46    /// The learning agent.
47    type LearningAgent: PolicyLearner<
48            TrainContext = Self::TrainingOutput,
49            InnerPolicy = Self::Policy,
50            Record: Checkpoint,
51        > + Send
52        + 'static;
53    /// The output data of a training step.
54    type TrainingOutput: ItemLazy + Clone + Send;
55}
56
57/// Concrete type that implements the [RLComponentsTypes](RLComponentsTypes) trait.
58pub struct RLComponentsMarker<E, EI, A> {
59    _env: PhantomData<E>,
60    _env_init: PhantomData<EI>,
61    _agent: PhantomData<A>,
62}
63
64impl<E, EI, A> RLComponentsTypes for RLComponentsMarker<E, EI, A>
65where
66    E: Environment + 'static,
67    EI: EnvironmentInit<E> + Send + 'static,
68    A: PolicyLearner + Send + 'static,
69    <A as PolicyLearner>::Record: Checkpoint,
70    A::TrainContext: ItemLazy + Clone + Send,
71    A::InnerPolicy: Policy + Send,
72    <A::InnerPolicy as Policy>::Observation: Batchable + Clone + Send,
73    <A::InnerPolicy as Policy>::ActionDistribution: Batchable + Clone + Send,
74    <A::InnerPolicy as Policy>::Action: Batchable + Clone + Send,
75    <A::InnerPolicy as Policy>::ActionContext: ItemLazy + Clone + Send + 'static,
76    <A::InnerPolicy as Policy>::PolicyState: Clone + Send,
77    <<A::InnerPolicy as Policy>::PolicyState as PolicyState>::Record: Checkpoint,
78    E::State: ToObservation<<A::InnerPolicy as Policy>::Observation> + Clone + Send + 'static,
79    E::Action: From<<A::InnerPolicy as Policy>::Action>
80        + ToAction<<A::InnerPolicy as Policy>::Action>
81        + Clone
82        + Send
83        + 'static,
84{
85    type Env = E;
86    type EnvInit = EI;
87    type LearningAgent = A;
88    type Policy = A::InnerPolicy;
89    type PolicyObs = <A::InnerPolicy as Policy>::Observation;
90    type PolicyAD = <A::InnerPolicy as Policy>::ActionDistribution;
91    type PolicyAction = <A::InnerPolicy as Policy>::Action;
92    type ActionContext = <A::InnerPolicy as Policy>::ActionContext;
93    type PolicyState = <A::InnerPolicy as Policy>::PolicyState;
94    type TrainingOutput = A::TrainContext;
95    type State = E::State;
96    type Action = E::Action;
97}
98
99pub(crate) type RlPolicy<RLC> =
100    <<RLC as RLComponentsTypes>::LearningAgent as PolicyLearner>::InnerPolicy;
101/// The event processor type for reinforcement learning.
102pub type RLEventProcessorType<RLC> = AsyncProcessorTraining<
103    RLEvent<<RLC as RLComponentsTypes>::TrainingOutput, <RLC as RLComponentsTypes>::ActionContext>,
104    AgentEvaluationEvent<<RLC as RLComponentsTypes>::ActionContext>,
105>;
106/// The record of the policy.
107pub type RLPolicyRecord<RLC> =
108    <<<RLC as RLComponentsTypes>::Policy as Policy>::PolicyState as PolicyState>::Record;
109/// The record of the learning agent.
110pub type RLAgentRecord<RLC> = <<RLC as RLComponentsTypes>::LearningAgent as PolicyLearner>::Record;
burn_train/learner/rl/components.rs

burn_train/learner/rl/
components.rs
