cais/simulation.py at main · trieschlab/cais · GitHub

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
import argparse
from matplotlib import pyplot as plt
import torch.distributions as D
from torch.optim.lr_scheduler import CosineAnnealingLR
from utils import *

if __name__ == '__main__':
    # parse simulation parameters
    parser = argparse.ArgumentParser()
    parser.add_argument('--method', default='cais', type=str, choices=['cais', 'controllability'])
    parser.add_argument('--mode', default='normal', type=str, choices=['normal', 'inverted', 'bimodal'])
    parser.add_argument('--n_taus', default=49, type=int)
    parser.add_argument('--surprise_weight', default=3., type=float)
    parser.add_argument('--lr', default=0.03, type=float)
    parser.add_argument('--smoothing', default=0.03, type=float)
    args = parser.parse_args()

    # simulation steps for 3 phases
    n_steps_total = 2000
    n_steps_phase1 = 200
    n_steps_phase2 = 1000

    # initialize taus (e.g. n_taus=49, taus={0.02, 0.04, ..., 0.98})
    taus = torch.linspace(1, args.n_taus, args.n_taus).view(1, -1) / (args.n_taus + 1)

    # initialize quantile values
    quantile_values = torch.linspace(-1, 1, args.n_taus).unsqueeze(0).repeat(9, 1).requires_grad_()

    # training components
    optimizer = torch.optim.Adam([quantile_values], lr=args.lr)
    scheduler = CosineAnnealingLR(optimizer, n_steps_total, eta_min=0)
    index_offset = torch.tensor([0, 2, 4, 6])

    # define outcome distributions according to mode
    match args.mode:
        case 'normal':
            baseline_outcome_distribution = D.Normal(0, 0.3)
            contingent_outcome_distribution = D.Normal(5, 0.3)
        case 'inverted':
            baseline_outcome_distribution = D.Normal(5, 0.3)
            contingent_outcome_distribution = D.Normal(0, 0.3)
        case 'bimodal':
            baseline_outcome_distribution = D.Normal(0, 0.3)
            mix = D.Categorical(torch.ones(2))
            comp = D.Normal(torch.tensor([-2., 2.]), torch.tensor([1., 1.]))
            contingent_outcome_distribution = D.MixtureSameFamily(mix, comp)

    # initialize simulation
    outcome = baseline_outcome_distribution.sample()
    action = torch.zeros(4, dtype=torch.long)
    motivation = None
    motivation_record = []

    # simulate
    for step in range(n_steps_total):
        # compute motivation
        match args.method:
            case 'cais':
                sense_of_agency = ((quantile_values[:-1] - quantile_values[-1:]).view(4, 2, -1) ** 2).detach().mean(-1)
            case 'controllability':
                entropy = compute_entropy(quantile_values.detach(), args.n_taus)
                sense_of_agency = entropy[-1] - entropy[:-1].view(4, 2)

        # to turn-off surprise: set surprise_weight to 0
        surprise = compute_surprise(quantile_values[:-1, :].detach(), action, outcome)
        temporal_motivation = sense_of_agency.mean(-1) + surprise * args.surprise_weight

        # exponential average
        if motivation is None:
            motivation = temporal_motivation
        else:
            motivation = temporal_motivation * args.smoothing + motivation * (1 - args.smoothing)
        motivation_record.append(motivation)

        # sample action based on motivation
        action_logits = torch.stack([motivation.mean().expand(4), motivation], dim=1)
        action = D.Categorical(logits=action_logits).sample()

        # get baseline outcome
        outcome = baseline_outcome_distribution.sample()

        # if in attached phase and contingent limb moves, override the outcome
        if n_steps_phase1 <= step <= n_steps_phase1 + n_steps_phase2:
            # assume the first entry is the controlling action
            if action[0].item() == 1:
                outcome = contingent_outcome_distribution.sample()

        # train the quantile regressor
        loss = compute_quantile_loss(torch.cat([quantile_values[:-1][action + index_offset], quantile_values[-1:]]), outcome, taus)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # scheduler.step()

    # example to plot motivation record
    record = torch.stack(motivation_record)
    for i in range(4):
        plt.plot(record[:, i])