repo = "bsuite"
export bsuite_params

const BSUITE_DATASETS = [
    "cartpole",
    "catch",
    "mountain_car"
]

types = [
    "full", 
    "full_train",
    "full_valid"
]

function bsuite_params()
    game = BSUITE_DATASETS
    shards = 0:4
    type = types

    @info game shards type
end

function bsuite_init()
    for env in BSUITE_DATASETS
        for stochasticity in 0.0:0.1:0.5
            for index in 0:4
                for type in types
                    register(
                        DataDep(
                            "rl-unplugged-$repo-$env-$stochasticity-$index-$type",
                            """
                            Dataset: RL Unplugged bsuite
                            Credits: https://arxiv.org/abs/2006.13888
                            Url: https://github.com/deepmind/deepmind-research/tree/master/rl_unplugged
                            Authors: Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine,
                            Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal,
                            Josh Merel, Daniel Mankowitz, Cosmin Paduraru, Gabriel
                            Dulac-Arnold, Jerry Li, Mohammad Norouzi, Matt Hoffman,
                            Ofir Nachum, George Tucker, Nicolas Heess, Nando deFreitas
                            Year: 2020

                            Data accompanying [RL Unplugged: Benchmarks for Offline Reinforcement Learning].
                            bsuite datasets are fairly light-weight and running experiments doesn't require 
                            too much compute. We recommend you to try bsuite, if you are interested in 
                            small-scale and easy to run offline RL datasets generated by stochastic environments
                            where the stochasticity of the environment is easy to control.
                            """,
                            "gs://rl_unplugged/$repo/$env/0_$stochasticity/$(index)_$type-00000-of-00001",
                            fetch_method = fetch_gc_file
                        )
                    )
                end
            end
        end
    end
end