Implement ValueDICE and some restructuring

2021-08-04 19:34:49 +02:00
parent cba42c6e4d
commit 5c40de66fa
8 changed files with 442 additions and 86 deletions
--- a/src/main.py
+++ b/src/main.py
@@ -47,8 +47,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
    if train:

        # make policy, train and test datasets, and send to 
-        policy = policy_class(config)
        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
+        # generate transform from train_dataset
+        transforms = generate_transforms(train_dataset)
+        policy = policy_class(config, transforms)
        cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
        train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)