├── README.md ├── examples ├── bc_config.py ├── finetune_and_eval.py ├── generate_data.py ├── id_config.py ├── point_mass_config.py └── pretrain.py ├── imitation_pretraining ├── __init__.py ├── algs │ ├── __init__.py │ ├── base.py │ ├── bc │ │ ├── __init__.py │ │ └── bc_learner.py │ ├── pretrained_models │ │ ├── __init__.py │ │ ├── r3m.py │ │ └── timm.py │ ├── registry.py │ └── rep_learning │ │ ├── __init__.py │ │ ├── contrastive_learner.py │ │ ├── inverse_dynamics_learner.py │ │ └── reconstruction_learner.py ├── data_utils │ ├── __init__.py │ ├── augmentations.py │ ├── batch_utils.py │ ├── observation_adapters.py │ ├── preprocess_data.py │ └── replay_buffer.py ├── envs │ ├── __init__.py │ ├── kitchen.py │ ├── metaworld.py │ ├── my_point_mass.xml │ ├── point_mass.py │ ├── registry.py │ └── wrappers.py ├── experiments │ ├── __init__.py │ ├── data_generation │ │ ├── __init__.py │ │ ├── common.py │ │ ├── kitchen.py │ │ ├── metaworld.py │ │ └── point_mass.py │ ├── evaluation │ │ ├── __init__.py │ │ ├── eval.py │ │ ├── loading.py │ │ └── visualization.py │ └── training │ │ ├── __init__.py │ │ ├── train.py │ │ ├── train_double_encoder.py │ │ └── training_utils.py └── networks │ ├── __init__.py │ ├── encoder.py │ ├── mlp.py │ ├── policy.py │ ├── registry.py │ ├── simple_conv.py │ └── spatial_layers.py ├── notebooks └── paper figures.ipynb ├── requirements.txt ├── setup.py └── sweeps ├── __init__.py ├── bc_finetune_fixed_encoder.py ├── bc_finetune_scratch_cpu.py ├── bc_finetune_scratch_gpu.py ├── configs ├── __init__.py ├── bc.py ├── config_utils.py ├── contrastive.py ├── generate_kitchen.py ├── generate_metaworld.py ├── generate_point_mass.py ├── inverse_dynamics.py ├── reconstruction.py └── test.py ├── data_configs.py ├── finetune_sweep_ft_data_size.py ├── finetune_sweep_pt_data_all.py ├── finetune_sweep_pt_data_size.py ├── finetune_sweep_pt_nstep.py ├── generate_kitchen.py ├── generate_metaworld_finetune.py ├── generate_metaworld_pretrain.py ├── generate_point_mass_finetune.py ├── generate_point_mass_pretrain.py ├── pretrain_alg_configs.py ├── pretrain_sweep_data_all.py ├── pretrain_sweep_data_size.py ├── pretrain_sweep_nstep.py ├── rep_compare_direct.py ├── sweep_utils.py └── test_train.py /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/README.md -------------------------------------------------------------------------------- /examples/bc_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/bc_config.py -------------------------------------------------------------------------------- /examples/finetune_and_eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/finetune_and_eval.py -------------------------------------------------------------------------------- /examples/generate_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/generate_data.py -------------------------------------------------------------------------------- /examples/id_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/id_config.py -------------------------------------------------------------------------------- /examples/point_mass_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/point_mass_config.py -------------------------------------------------------------------------------- /examples/pretrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/examples/pretrain.py -------------------------------------------------------------------------------- /imitation_pretraining/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/algs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/__init__.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/base.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/bc/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/bc/__init__.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/bc/bc_learner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/bc/bc_learner.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/pretrained_models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/algs/pretrained_models/r3m.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/pretrained_models/r3m.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/pretrained_models/timm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/pretrained_models/timm.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/registry.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/registry.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/rep_learning/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /imitation_pretraining/algs/rep_learning/contrastive_learner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/rep_learning/contrastive_learner.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/rep_learning/inverse_dynamics_learner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/rep_learning/inverse_dynamics_learner.py -------------------------------------------------------------------------------- /imitation_pretraining/algs/rep_learning/reconstruction_learner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/algs/rep_learning/reconstruction_learner.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/__init__.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/augmentations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/augmentations.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/batch_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/batch_utils.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/observation_adapters.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/observation_adapters.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/preprocess_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/preprocess_data.py -------------------------------------------------------------------------------- /imitation_pretraining/data_utils/replay_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/data_utils/replay_buffer.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/__init__.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/kitchen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/kitchen.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/metaworld.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/metaworld.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/my_point_mass.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/my_point_mass.xml -------------------------------------------------------------------------------- /imitation_pretraining/envs/point_mass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/point_mass.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/registry.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/registry.py -------------------------------------------------------------------------------- /imitation_pretraining/envs/wrappers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/envs/wrappers.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/experiments/data_generation/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/experiments/data_generation/common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/data_generation/common.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/data_generation/kitchen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/data_generation/kitchen.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/data_generation/metaworld.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/data_generation/metaworld.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/data_generation/point_mass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/data_generation/point_mass.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/experiments/evaluation/eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/evaluation/eval.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/evaluation/loading.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/evaluation/loading.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/evaluation/visualization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/evaluation/visualization.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/training/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imitation_pretraining/experiments/training/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/training/train.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/training/train_double_encoder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/training/train_double_encoder.py -------------------------------------------------------------------------------- /imitation_pretraining/experiments/training/training_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/experiments/training/training_utils.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/__init__.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/encoder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/encoder.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/mlp.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/policy.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/registry.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/registry.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/simple_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/simple_conv.py -------------------------------------------------------------------------------- /imitation_pretraining/networks/spatial_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/imitation_pretraining/networks/spatial_layers.py -------------------------------------------------------------------------------- /notebooks/paper figures.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/notebooks/paper figures.ipynb -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/requirements.txt -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/setup.py -------------------------------------------------------------------------------- /sweeps/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /sweeps/bc_finetune_fixed_encoder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/bc_finetune_fixed_encoder.py -------------------------------------------------------------------------------- /sweeps/bc_finetune_scratch_cpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/bc_finetune_scratch_cpu.py -------------------------------------------------------------------------------- /sweeps/bc_finetune_scratch_gpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/bc_finetune_scratch_gpu.py -------------------------------------------------------------------------------- /sweeps/configs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/__init__.py -------------------------------------------------------------------------------- /sweeps/configs/bc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/bc.py -------------------------------------------------------------------------------- /sweeps/configs/config_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/config_utils.py -------------------------------------------------------------------------------- /sweeps/configs/contrastive.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/contrastive.py -------------------------------------------------------------------------------- /sweeps/configs/generate_kitchen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/generate_kitchen.py -------------------------------------------------------------------------------- /sweeps/configs/generate_metaworld.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/generate_metaworld.py -------------------------------------------------------------------------------- /sweeps/configs/generate_point_mass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/generate_point_mass.py -------------------------------------------------------------------------------- /sweeps/configs/inverse_dynamics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/inverse_dynamics.py -------------------------------------------------------------------------------- /sweeps/configs/reconstruction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/reconstruction.py -------------------------------------------------------------------------------- /sweeps/configs/test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/configs/test.py -------------------------------------------------------------------------------- /sweeps/data_configs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/data_configs.py -------------------------------------------------------------------------------- /sweeps/finetune_sweep_ft_data_size.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/finetune_sweep_ft_data_size.py -------------------------------------------------------------------------------- /sweeps/finetune_sweep_pt_data_all.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/finetune_sweep_pt_data_all.py -------------------------------------------------------------------------------- /sweeps/finetune_sweep_pt_data_size.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/finetune_sweep_pt_data_size.py -------------------------------------------------------------------------------- /sweeps/finetune_sweep_pt_nstep.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/finetune_sweep_pt_nstep.py -------------------------------------------------------------------------------- /sweeps/generate_kitchen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/generate_kitchen.py -------------------------------------------------------------------------------- /sweeps/generate_metaworld_finetune.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/generate_metaworld_finetune.py -------------------------------------------------------------------------------- /sweeps/generate_metaworld_pretrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/generate_metaworld_pretrain.py -------------------------------------------------------------------------------- /sweeps/generate_point_mass_finetune.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/generate_point_mass_finetune.py -------------------------------------------------------------------------------- /sweeps/generate_point_mass_pretrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/generate_point_mass_pretrain.py -------------------------------------------------------------------------------- /sweeps/pretrain_alg_configs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/pretrain_alg_configs.py -------------------------------------------------------------------------------- /sweeps/pretrain_sweep_data_all.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/pretrain_sweep_data_all.py -------------------------------------------------------------------------------- /sweeps/pretrain_sweep_data_size.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/pretrain_sweep_data_size.py -------------------------------------------------------------------------------- /sweeps/pretrain_sweep_nstep.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/pretrain_sweep_nstep.py -------------------------------------------------------------------------------- /sweeps/rep_compare_direct.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/rep_compare_direct.py -------------------------------------------------------------------------------- /sweeps/sweep_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/sweep_utils.py -------------------------------------------------------------------------------- /sweeps/test_train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/davidbrandfonbrener/imitation_pretraining/HEAD/sweeps/test_train.py --------------------------------------------------------------------------------