├── .gitignore ├── Readme.md ├── __init__.py ├── configs └── atari │ ├── finetune │ ├── agent │ │ ├── der.yaml │ │ ├── drq.yaml │ │ └── drq_freeze.yaml │ ├── drq.yaml │ ├── env │ │ └── atari.yaml │ ├── model │ │ ├── backbone │ │ │ └── impala.yaml │ │ ├── head │ │ │ └── identity.yaml │ │ ├── impala.yaml │ │ └── policy │ │ │ └── rainbow_d512_std05.yaml │ ├── pretrain │ │ ├── scratch.yaml │ │ └── simtpr.yaml │ └── simtpr.yaml │ └── pretrain │ ├── dataloader │ ├── act │ │ ├── r2_ckpt50_n64_t1_f4.yaml │ │ └── r2_ckpt50_n64_t4_f1.yaml │ ├── lightweight_gpu.yaml │ ├── rew │ │ ├── r2_ckpt1_n64_t1_f4.yaml │ │ ├── r2_ckpt1_n64_t4_f1.yaml │ │ └── r2_ckpt50_n64_t1_f4.yaml │ └── train │ │ ├── r1_ckpt345_n320_t2_f4.yaml │ │ ├── r1_ckpt345_n32_t11_f4.yaml │ │ ├── r1_ckpt345_n640_t1_f4.yaml │ │ ├── r1_ckpt345_n64_t10_f4.yaml │ │ ├── r1_ckpt345_n64_t11_f4.yaml │ │ └── r1_ckpt345_n64_t13_f4.yaml │ ├── env │ └── atari.yaml │ ├── model │ ├── backbone │ │ └── impala.yaml │ ├── head │ │ └── simtpr.yaml │ ├── policy │ │ └── identity.yaml │ └── simtpr.yaml │ ├── pretrain │ └── scratch.yaml │ ├── simtpr.yaml │ └── trainer │ └── simtpr.yaml ├── data └── download_atari_replay_dataset.sh ├── download_models.py ├── requirements.yml ├── run_atari_finetune.py ├── run_finetune.py ├── run_pretrain.py └── src ├── __init__.py ├── agents ├── __init__.py ├── base.py ├── buffers │ ├── __init__.py │ ├── base.py │ └── per_buffer.py └── rainbow.py ├── common ├── __init__.py ├── augmentation.py ├── class_utils.py ├── data_utils.py ├── logger.py ├── losses.py ├── train_utils.py ├── vis_utils.py └── vit_utils.py ├── dataloaders ├── __init__.py ├── base.py └── replay.py ├── envs ├── __init__.py ├── atari.py ├── base.py └── vec_env.py ├── models ├── __init__.py ├── backbones │ ├── __init__.py │ ├── base.py │ └── cnn │ │ ├── __init__.py │ │ ├── de_nature.py │ │ ├── impala.py │ │ └── nature.py ├── base.py ├── heads │ ├── __init__.py │ ├── base.py │ ├── identity_head.py │ └── simtpr_head.py ├── layers.py └── policies │ ├── __init__.py │ ├── base.py │ ├── identity_policy.py │ └── rainbow_policy.py └── trainers ├── __init__.py ├── base.py └── simtpr.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/.gitignore -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/Readme.md -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /configs/atari/finetune/agent/der.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/agent/der.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/agent/drq.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/agent/drq.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/agent/drq_freeze.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/agent/drq_freeze.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/drq.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/drq.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/env/atari.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/env/atari.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/model/backbone/impala.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/model/backbone/impala.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/model/head/identity.yaml: -------------------------------------------------------------------------------- 1 | type: identity 2 | in_dim: None -------------------------------------------------------------------------------- /configs/atari/finetune/model/impala.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/model/impala.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/model/policy/rainbow_d512_std05.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/model/policy/rainbow_d512_std05.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/pretrain/scratch.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/pretrain/scratch.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/pretrain/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/pretrain/simtpr.yaml -------------------------------------------------------------------------------- /configs/atari/finetune/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/finetune/simtpr.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/act/r2_ckpt50_n64_t1_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/act/r2_ckpt50_n64_t1_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/act/r2_ckpt50_n64_t4_f1.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/act/r2_ckpt50_n64_t4_f1.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/lightweight_gpu.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/lightweight_gpu.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/rew/r2_ckpt1_n64_t1_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/rew/r2_ckpt1_n64_t1_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/rew/r2_ckpt1_n64_t4_f1.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/rew/r2_ckpt1_n64_t4_f1.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/rew/r2_ckpt50_n64_t1_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/rew/r2_ckpt50_n64_t1_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n320_t2_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n320_t2_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n32_t11_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n32_t11_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n640_t1_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n640_t1_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t10_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t10_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t11_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t11_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t13_f4.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/dataloader/train/r1_ckpt345_n64_t13_f4.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/env/atari.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/env/atari.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/model/backbone/impala.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/model/backbone/impala.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/model/head/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/model/head/simtpr.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/model/policy/identity.yaml: -------------------------------------------------------------------------------- 1 | type: identity 2 | in_dim: None 3 | -------------------------------------------------------------------------------- /configs/atari/pretrain/model/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/model/simtpr.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/pretrain/scratch.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/pretrain/scratch.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/simtpr.yaml -------------------------------------------------------------------------------- /configs/atari/pretrain/trainer/simtpr.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/configs/atari/pretrain/trainer/simtpr.yaml -------------------------------------------------------------------------------- /data/download_atari_replay_dataset.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/data/download_atari_replay_dataset.sh -------------------------------------------------------------------------------- /download_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/download_models.py -------------------------------------------------------------------------------- /requirements.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/requirements.yml -------------------------------------------------------------------------------- /run_atari_finetune.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/run_atari_finetune.py -------------------------------------------------------------------------------- /run_finetune.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/run_finetune.py -------------------------------------------------------------------------------- /run_pretrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/run_pretrain.py -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/agents/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/__init__.py -------------------------------------------------------------------------------- /src/agents/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/base.py -------------------------------------------------------------------------------- /src/agents/buffers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/buffers/__init__.py -------------------------------------------------------------------------------- /src/agents/buffers/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/buffers/base.py -------------------------------------------------------------------------------- /src/agents/buffers/per_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/buffers/per_buffer.py -------------------------------------------------------------------------------- /src/agents/rainbow.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/agents/rainbow.py -------------------------------------------------------------------------------- /src/common/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/common/augmentation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/augmentation.py -------------------------------------------------------------------------------- /src/common/class_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/class_utils.py -------------------------------------------------------------------------------- /src/common/data_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/data_utils.py -------------------------------------------------------------------------------- /src/common/logger.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/logger.py -------------------------------------------------------------------------------- /src/common/losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/losses.py -------------------------------------------------------------------------------- /src/common/train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/train_utils.py -------------------------------------------------------------------------------- /src/common/vis_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/vis_utils.py -------------------------------------------------------------------------------- /src/common/vit_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/common/vit_utils.py -------------------------------------------------------------------------------- /src/dataloaders/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/dataloaders/__init__.py -------------------------------------------------------------------------------- /src/dataloaders/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/dataloaders/base.py -------------------------------------------------------------------------------- /src/dataloaders/replay.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/dataloaders/replay.py -------------------------------------------------------------------------------- /src/envs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/envs/__init__.py -------------------------------------------------------------------------------- /src/envs/atari.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/envs/atari.py -------------------------------------------------------------------------------- /src/envs/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/envs/base.py -------------------------------------------------------------------------------- /src/envs/vec_env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/envs/vec_env.py -------------------------------------------------------------------------------- /src/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/__init__.py -------------------------------------------------------------------------------- /src/models/backbones/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/backbones/__init__.py -------------------------------------------------------------------------------- /src/models/backbones/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/backbones/base.py -------------------------------------------------------------------------------- /src/models/backbones/cnn/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/backbones/cnn/de_nature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/backbones/cnn/de_nature.py -------------------------------------------------------------------------------- /src/models/backbones/cnn/impala.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/backbones/cnn/impala.py -------------------------------------------------------------------------------- /src/models/backbones/cnn/nature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/backbones/cnn/nature.py -------------------------------------------------------------------------------- /src/models/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/base.py -------------------------------------------------------------------------------- /src/models/heads/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/heads/__init__.py -------------------------------------------------------------------------------- /src/models/heads/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/heads/base.py -------------------------------------------------------------------------------- /src/models/heads/identity_head.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/heads/identity_head.py -------------------------------------------------------------------------------- /src/models/heads/simtpr_head.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/heads/simtpr_head.py -------------------------------------------------------------------------------- /src/models/layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/layers.py -------------------------------------------------------------------------------- /src/models/policies/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/policies/__init__.py -------------------------------------------------------------------------------- /src/models/policies/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/policies/base.py -------------------------------------------------------------------------------- /src/models/policies/identity_policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/policies/identity_policy.py -------------------------------------------------------------------------------- /src/models/policies/rainbow_policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/models/policies/rainbow_policy.py -------------------------------------------------------------------------------- /src/trainers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/trainers/__init__.py -------------------------------------------------------------------------------- /src/trainers/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/trainers/base.py -------------------------------------------------------------------------------- /src/trainers/simtpr.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dojeon-ai/SimTPR/HEAD/src/trainers/simtpr.py --------------------------------------------------------------------------------