InContextContinualLearning/td_data_generation.py at main · Sequential-Intelligence-Lab/InContextContinualLearning · 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
import os
import pickle
from argparse import ArgumentParser, Namespace

import numpy as np
from tqdm import tqdm

from rl.feat import Feature
from rl.MRP.boyan import BoyanChain
from rl.td_task import TDTask
from utils import set_seed


def gen_td_data(n_states: int, gamma: float,
                d_feat: int, feat_mode: str, representable: bool,
                num_tasks: int, steps_per_task: int, seed: int):
    set_seed(seed)
    tasks = []
    for _ in tqdm(range(num_tasks)):
        feat = Feature(d=d_feat, s=n_states, mode=feat_mode)
        if representable:
            w = np.randm.randn(d_feat, 1)
            bc = BoyanChain(n_states=n_states, gamma=gamma,
                            weight=w, phi=feat.phi)
        else:
            bc = BoyanChain(n_states=n_states, gamma=gamma)
        s_arr = np.zeros(steps_per_task, dtype=np.int32)
        r_arr = np.zeros(steps_per_task, dtype=np.float32)
        sp_arr = np.zeros(steps_per_task, dtype=np.int32)

        s = bc.reset()
        for i in range(steps_per_task):
            sp, r = bc.step(s)
            s_arr[i] = s
            r_arr[i] = r
            sp_arr[i] = sp
            s = sp
        task = TDTask(mrp=bc, feat=feat, s=s_arr, r=r_arr, sp=sp_arr)
        tasks.append(task)

    dir_path = os.path.join('.', 'data', 'td')
    os.makedirs(dir_path, exist_ok=True)
    with open(os.path.join(dir_path, f'{seed}.pkl'), 'wb') as f:
        pickle.dump(tasks, f)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-s', '--num_states', type=int,
                        help='number of states', default=10)
    parser.add_argument('--gamma', type=float,
                        help='discount factor', default=0.9)
    parser.add_argument('--dim_feature', type=int,
                        help='feature dimension', default=4)
    parser.add_argument('--feat_mode', type=str,
                        help='feature mode', default='random', choices=['random', 'one-hot'])
    parser.add_argument('--representable', action='store_true',
                        help='representable value function')
    parser.add_argument('--num_tasks', type=int,
                        help='number of tasks', default=50_000)
    parser.add_argument('--steps_per_task', type=int,
                        help='number of interaction steps per task', default=500)
    parser.add_argument('--seed', type=int, nargs='+',
                        help='random seed', default=list(range(10)))

    args: Namespace = parser.parse_args()

    print('Generating TD data...')

    for seed in args.seed:
        config = dict(n_states=args.num_states,
                      gamma=args.gamma,
                      d_feat=args.dim_feature,
                      feat_mode=args.feat_mode,
                      representable=args.representable,
                      num_tasks=args.num_tasks,
                      steps_per_task=args.steps_per_task,
                      seed=seed)
        gen_td_data(**config)