rsrl 0.8.1

A fast, extensible reinforcement learning framework in Rust
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

use crate::{
    domains::Transition,
    fa::ScaledGradientUpdate,
    params::BufferMut,
    Differentiable,
    Handler,
};

#[derive(Clone, Debug)]
#[cfg_attr(
    feature = "serde",
    derive(Serialize, Deserialize),
    serde(crate = "serde_crate")
)]
pub struct Response<T, W> {
    pub res_theta: T,
    pub res_w: W,
    pub td_error: f64,
}

#[derive(Clone, Debug, Parameterised)]
#[cfg_attr(
    feature = "serde",
    derive(Serialize, Deserialize),
    serde(crate = "serde_crate")
)]
pub struct GTD2<F> {
    #[weights]
    pub fa_theta: F,
    pub fa_w: F,

    pub gamma: f64,
}

type SGU<'m, S, F> = ScaledGradientUpdate<<F as Differentiable<(&'m S,)>>::Jacobian>;
type SGURef<'m, 'j, S, F> = ScaledGradientUpdate<&'j <F as Differentiable<(&'m S,)>>::Jacobian>;

impl<'m, S, A, F> Handler<&'m Transition<S, A>> for GTD2<F>
where
    F: Differentiable<(&'m S,), Output = f64> + Handler<SGU<'m, S, F>>,
    F: for<'j> Handler<
        SGURef<'m, 'j, S, F>,
        Response = <F as Handler<SGU<'m, S, F>>>::Response,
        Error = <F as Handler<SGU<'m, S, F>>>::Error,
    >,
{
    type Response =
        Response<<F as Handler<SGU<'m, S, F>>>::Response, <F as Handler<SGU<'m, S, F>>>::Response>;
    type Error = <F as Handler<SGU<'m, S, F>>>::Error;

    fn handle(&mut self, t: &'m Transition<S, A>) -> Result<Self::Response, Self::Error> {
        let (s, ns) = t.states();

        let w_s = self.fa_w.evaluate((s,));
        let theta_s = self.fa_theta.evaluate((s,));
        let theta_ns = self.fa_theta.evaluate((ns,));

        let td_error = if t.terminated() {
            t.reward - theta_s
        } else {
            t.reward + self.gamma * theta_ns - theta_s
        };

        let mut grad_s = self.fa_theta.grad((s,));

        let res_w = self.fa_w.handle(ScaledGradientUpdate {
            alpha: td_error - w_s,
            jacobian: &grad_s,
        })?;

        let grad_ns = self.fa_theta.grad((ns,));

        grad_s.merge_inplace(&grad_ns, |x, y| x - self.gamma * y);

        let res_theta = self.fa_theta.handle(ScaledGradientUpdate {
            alpha: w_s,
            jacobian: grad_s,
        })?;

        Ok(Response {
            res_theta,
            res_w,
            td_error,
        })
    }
}