Policy Gradient Theorem

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Simple 1D Random Walk environment.
5
// States: 0 ... 4 (0 and 4 are terminal). Start at 2.
6
// Actions: 0 = Left, 1 = Right. Rewards: +1 at state 4, -1 at state 0, otherwise 0.
7
struct RandomWalkEnv {
8
    int n_states = 5;
9
    int start = 2;
10
    int s;
11
    RandomWalkEnv() { reset(); }
12
    int reset() { s = start; return s; }
13
    // Step returns: next_state, reward, done
14
    tuple<int, double, bool> step(int a) {
15
        if (s == 0 || s == 4) return {s, 0.0, true};
16
        int ns = s + (a == 1 ? 1 : -1);
17
        ns = max(0, min(4, ns));
18
        double r = (ns == 4 ? 1.0 : (ns == 0 ? -1.0 : 0.0));
19
        bool done = (ns == 0 || ns == 4);
20
        s = ns;
21
        return {ns, r, done};
22
    }
23
};
24
25
// Softmax policy over 2 actions with tabular parameters theta[s][a].
26
struct SoftmaxPolicy {
27
    int S, A;
28
    vector<array<double,2>> theta; // parameters per state and action
29
    SoftmaxPolicy(int S_, int A_=2) : S(S_), A(A_) {
30
        theta.assign(S, {0.0, 0.0});
31
    }
32
    array<double,2> probs(int s) const {
33
        double m = max(theta[s][0], theta[s][1]); // for numerical stability
34
        double e0 = exp(theta[s][0] - m);
35
        double e1 = exp(theta[s][1] - m);
36
        double Z = e0 + e1;
37
        return { e0 / Z, e1 / Z };
38
    }
39
    int sample_action(int s, mt19937 &gen) const {
40
        auto p = probs(s);
41
        uniform_real_distribution<> dist(0.0, 1.0);
42
        double u = dist(gen);
43
        return (u < p[0]) ? 0 : 1;
44
    }
45
};
46
47
int main() {
48
    ios::sync_with_stdio(false);
49
    cin.tie(nullptr);
50
51
    RandomWalkEnv env;
52
    int S = env.n_states, A = 2;
53
    SoftmaxPolicy pi(S, A);
54
55
    // Baseline (state-value) V(s)
56
    vector<double> V(S, 0.0);
57
58
    // Hyperparameters
59
    double gamma = 0.99;              // discount
60
    double alpha_pi = 0.05;           // policy step size
61
    double alpha_v = 0.05;            // value step size
62
    int max_episodes = 2000;
63
    int max_steps = 30;               // safety cap on episode length
64
    mt19937 gen(42);
65
66
    for (int ep = 0; ep < max_episodes; ++ep) {
67
        // Generate one episode
68
        vector<int> states, actions;
69
        vector<double> rewards;
70
        int s = env.reset();
71
        bool done = false;
72
        for (int t = 0; t < max_steps && !done; ++t) {
73
            int a = pi.sample_action(s, gen);
74
            auto [ns, r, d] = env.step(a);
75
            states.push_back(s);
76
            actions.push_back(a);
77
            rewards.push_back(r);
78
            s = ns; done = d;
79
        }
80
        int T = (int)rewards.size();
81
82
        // Compute discounted returns G_t
83
        vector<double> G(T, 0.0);
84
        double g = 0.0;
85
        for (int t = T - 1; t >= 0; --t) {
86
            g = rewards[t] + gamma * g;
87
            G[t] = g;
88
        }
89
90
        // Policy and value updates using advantage A_t = G_t - V(s_t)
91
        for (int t = 0; t < T; ++t) {
92
            int st = states[t], at = actions[t];
93
            auto p = pi.probs(st);
94
            double Ahat = G[t] - V[st]; // baseline reduces variance
95
96
            // Gradient of log-softmax: d/dtheta(s,a') log pi(a|s) = 1_{a'=a} - pi(a'|s)
97
            for (int a = 0; a < A; ++a) {
98
                double grad_logpi = (a == at ? 1.0 : 0.0) - (a == 0 ? p[0] : p[1]);
99
                pi.theta[st][a] += alpha_pi * grad_logpi * Ahat;
100
            }
101
            // Update baseline towards return (value regression)
102
            V[st] += alpha_v * (G[t] - V[st]);
103
        }
104
105
        // Optional: monitor learning every 200 episodes
106
        if ((ep + 1) % 200 == 0) {
107
            // Evaluate greedy tendency from center state
108
            auto p = pi.probs(2);
109
            cout << "Episode " << (ep + 1)
110
                 << ": P(Left|s=2)=" << fixed << setprecision(3) << p[0]
111
                 << ", P(Right|s=2)=" << p[1] << "\n";
112
        }
113
    }
114
115
    // Final policy probabilities for all states
116
    cout << "Final policy (P(Left), P(Right)) per state:\n";
117
    for (int s = 0; s < S; ++s) {
118
        auto p = pi.probs(s);
119
        cout << "s=" << s << ": (" << fixed << setprecision(3) << p[0]
120
             << ", " << p[1] << ")\n";
121
    }
122
    return 0;
123
}
124

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
struct RandomWalkEnv {
5
    int n_states = 5;
6
    int start = 2;
7
    int s;
8
    RandomWalkEnv() { reset(); }
9
    int reset() { s = start; return s; }
10
    tuple<int, double, bool> step(int a) {
11
        if (s == 0 || s == 4) return {s, 0.0, true};
12
        int ns = s + (a == 1 ? 1 : -1);
13
        ns = max(0, min(4, ns));
14
        double r = (ns == 4 ? 1.0 : (ns == 0 ? -1.0 : 0.0));
15
        bool done = (ns == 0 || ns == 4);
16
        s = ns;
17
        return {ns, r, done};
18
    }
19
};
20
21
struct SoftmaxPolicy {
22
    int S, A;
23
    vector<array<double,2>> theta;
24
    SoftmaxPolicy(int S_, int A_=2) : S(S_), A(A_) { theta.assign(S, {0.0, 0.0}); }
25
    array<double,2> probs(int s) const {
26
        double m = max(theta[s][0], theta[s][1]);
27
        double e0 = exp(theta[s][0] - m);
28
        double e1 = exp(theta[s][1] - m);
29
        double Z = e0 + e1;
30
        return { e0 / Z, e1 / Z };
31
    }
32
    int sample_action(int s, mt19937 &gen) const {
33
        auto p = probs(s);
34
        uniform_real_distribution<> dist(0.0, 1.0);
35
        double u = dist(gen);
36
        return (u < p[0]) ? 0 : 1;
37
    }
38
};
39
40
int main(){
41
    ios::sync_with_stdio(false);
42
    cin.tie(nullptr);
43
44
    RandomWalkEnv env;
45
    int S = env.n_states, A = 2;
46
    SoftmaxPolicy pi(S, A);
47
    vector<double> V(S, 0.0);
48
49
    double gamma = 0.99;
50
    double alpha_pi = 0.05;  // actor step size
51
    double alpha_v  = 0.1;   // critic step size (often larger)
52
    int episodes = 3000;
53
    int max_steps = 30;
54
    mt19937 gen(123);
55
56
    for(int ep=0; ep<episodes; ++ep){
57
        int s = env.reset();
58
        bool done = false;
59
        for(int t=0; t<max_steps && !done; ++t){
60
            // Sample action
61
            int a = pi.sample_action(s, gen);
62
            auto p = pi.probs(s);
63
            auto [ns, r, d] = env.step(a);
64
65
            // TD error: delta = r + gamma * V(ns) - V(s)
66
            double delta = r + (d ? 0.0 : gamma * V[ns]) - V[s];
67
68
            // Critic update
69
            V[s] += alpha_v * delta;
70
71
            // Actor update using TD error as advantage estimate
72
            for(int ap=0; ap<A; ++ap){
73
                double grad_logpi = (ap == a ? 1.0 : 0.0) - (ap == 0 ? p[0] : p[1]);
74
                pi.theta[s][ap] += alpha_pi * grad_logpi * delta; // online actor-critic
75
            }
76
77
            s = ns; done = d;
78
        }
79
        if((ep+1) % 300 == 0){
80
            auto p2 = pi.probs(2);
81
            cout << "Ep " << (ep+1) << ": P(Right|2)=" << fixed << setprecision(3) << p2[1] << "\n";
82
        }
83
    }
84
85
    cout << "Learned policy probs at center state s=2: ";
86
    auto p2 = pi.probs(2);
87
    cout << "(" << fixed << setprecision(3) << p2[0] << ", " << p2[1] << ")\n";
88
    return 0;
89
}
90

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Two-armed stochastic bandit with Gaussian rewards.
5
struct Bandit2 {
6
    double mu0, mu1;  // mean rewards for arms 0 and 1
7
    mt19937 gen;
8
    normal_distribution<> noise;
9
    Bandit2(double m0, double m1, int seed=7) : mu0(m0), mu1(m1), gen(seed), noise(0.0, 1.0) {}
10
    // Pull arm a -> reward ~ N(mu_a, 1)
11
    double pull(int a){ return (a==0 ? mu0 : mu1) + noise(gen); }
12
};
13
14
struct Softmax2 {
15
    array<double,2> theta{0.0, 0.0};
16
    array<double,2> probs() const {
17
        double m = max(theta[0], theta[1]);
18
        double e0 = exp(theta[0]-m), e1 = exp(theta[1]-m);
19
        double Z = e0 + e1; return {e0/Z, e1/Z};
20
    }
21
    int sample(mt19937 &g) const {
22
        auto p = probs();
23
        uniform_real_distribution<> U(0.0,1.0);
24
        return (U(g) < p[0]) ? 0 : 1;
25
    }
26
};
27
28
int main(){
29
    ios::sync_with_stdio(false); cin.tie(nullptr);
30
    Bandit2 env(0.0, 0.5);      // arm 1 is better
31
    Softmax2 pi;
32
    mt19937 gen(1234);
33
34
    double alpha = 0.1;         // natural gradient step size
35
    double baseline = 0.0;      // running average reward as baseline
36
    double beta = 0.01;         // baseline update rate
37
38
    for(int t=1; t<=5000; ++t){
39
        auto p = pi.probs();
40
        int a = pi.sample(gen);
41
        double r = env.pull(a);
42
        baseline += beta * (r - baseline);
43
44
        // Score function gradient g = (r - b) * (onehot(a) - p)
45
        array<double,2> onehot{ (a==0)?1.0:0.0, (a==1)?1.0:0.0 };
46
        array<double,2> g{
47
            (r - baseline) * (onehot[0] - p[0]),
48
            (r - baseline) * (onehot[1] - p[1])
49
        };
50
51
        // Fisher for 2-action softmax: F = diag(p) - p p^T
52
        // Compute 2x2 F and its inverse explicitly.
53
        double p0 = p[0], p1 = p[1];
54
        double F00 = p0*(1.0 - p0);
55
        double F11 = p1*(1.0 - p1);
56
        double F01 = -p0*p1;
57
        double F10 = -p0*p1;
58
        // Inverse of 2x2 matrix [[a,b],[c,d]] is (1/det) * [[d,-b],[-c,a]]
59
        double det = F00*F11 - F01*F10;
60
        // Guard against near-singularity when probabilities saturate
61
        if (det < 1e-8) det = 1e-8;
62
        double iF00 =  F11 / det;
63
        double iF01 = -F01 / det;
64
        double iF10 = -F10 / det;
65
        double iF11 =  F00 / det;
66
67
        // Natural gradient: theta += alpha * F^{-1} * g
68
        array<double,2> ng{
69
            iF00 * g[0] + iF01 * g[1],
70
            iF10 * g[0] + iF11 * g[1]
71
        };
72
        pi.theta[0] += alpha * ng[0];
73
        pi.theta[1] += alpha * ng[1];
74
75
        if (t % 500 == 0) {
76
            auto q = pi.probs();
77
            cout << "t=" << t << ", P(arm1)=" << fixed << setprecision(3) << q[1]
78
                 << ", avgR~" << baseline << "\n";
79
        }
80
    }
81
82
    auto p = pi.probs();
83
    cout << "Final probabilities: P(arm0)=" << fixed << setprecision(3) << p[0]
84
         << ", P(arm1)=" << p[1] << "\n";
85
    return 0;
86
}
87