├── .gitignore ├── README.md ├── media ├── cheetah-run.png ├── finger-turn-hard.png ├── humanoid-stand.png └── walker-run.png ├── setup.py └── tdmpc2_jax ├── __init__.py ├── common ├── __init__.py ├── activations.py ├── loss.py ├── scale.py └── util.py ├── config.yaml ├── data ├── __init__.py ├── episodic_buffer.py └── sequential_buffer.py ├── envs ├── __init__.py ├── dmcontrol.py ├── tasks │ ├── __init__.py │ ├── ball_in_cup.py │ ├── ball_in_cup.xml │ ├── cheetah.py │ ├── cheetah.xml │ ├── fish.py │ ├── fish.xml │ ├── hopper.py │ ├── hopper.xml │ ├── pendulum.py │ ├── pendulum.xml │ ├── reacher.py │ ├── reacher_four_links.xml │ ├── reacher_three_links.xml │ ├── walker.py │ └── walker.xml └── wrappers │ ├── __init__.py │ ├── action_repeat.py │ ├── pixels.py │ └── time_limit.py ├── networks ├── __init__.py ├── ensemble.py └── mlp.py ├── tdmpc2.py ├── train.py └── world_model.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/.gitignore -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/README.md -------------------------------------------------------------------------------- /media/cheetah-run.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/media/cheetah-run.png -------------------------------------------------------------------------------- /media/finger-turn-hard.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/media/finger-turn-hard.png -------------------------------------------------------------------------------- /media/humanoid-stand.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/media/humanoid-stand.png -------------------------------------------------------------------------------- /media/walker-run.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/media/walker-run.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/setup.py -------------------------------------------------------------------------------- /tdmpc2_jax/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/__init__.py -------------------------------------------------------------------------------- /tdmpc2_jax/common/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tdmpc2_jax/common/activations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/common/activations.py -------------------------------------------------------------------------------- /tdmpc2_jax/common/loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/common/loss.py -------------------------------------------------------------------------------- /tdmpc2_jax/common/scale.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/common/scale.py -------------------------------------------------------------------------------- /tdmpc2_jax/common/util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/common/util.py -------------------------------------------------------------------------------- /tdmpc2_jax/config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/config.yaml -------------------------------------------------------------------------------- /tdmpc2_jax/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/data/__init__.py -------------------------------------------------------------------------------- /tdmpc2_jax/data/episodic_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/data/episodic_buffer.py -------------------------------------------------------------------------------- /tdmpc2_jax/data/sequential_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/data/sequential_buffer.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/__init__.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/dmcontrol.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/dmcontrol.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/ball_in_cup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/ball_in_cup.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/ball_in_cup.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/ball_in_cup.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/cheetah.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/cheetah.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/cheetah.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/cheetah.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/fish.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/fish.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/fish.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/fish.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/hopper.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/hopper.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/hopper.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/hopper.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/pendulum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/pendulum.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/pendulum.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/pendulum.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/reacher.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/reacher.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/reacher_four_links.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/reacher_four_links.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/reacher_three_links.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/reacher_three_links.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/walker.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/walker.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/tasks/walker.xml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/tasks/walker.xml -------------------------------------------------------------------------------- /tdmpc2_jax/envs/wrappers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tdmpc2_jax/envs/wrappers/action_repeat.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/wrappers/action_repeat.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/wrappers/pixels.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/wrappers/pixels.py -------------------------------------------------------------------------------- /tdmpc2_jax/envs/wrappers/time_limit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/envs/wrappers/time_limit.py -------------------------------------------------------------------------------- /tdmpc2_jax/networks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/networks/__init__.py -------------------------------------------------------------------------------- /tdmpc2_jax/networks/ensemble.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/networks/ensemble.py -------------------------------------------------------------------------------- /tdmpc2_jax/networks/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/networks/mlp.py -------------------------------------------------------------------------------- /tdmpc2_jax/tdmpc2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/tdmpc2.py -------------------------------------------------------------------------------- /tdmpc2_jax/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/train.py -------------------------------------------------------------------------------- /tdmpc2_jax/world_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShaneFlandermeyer/tdmpc2-jax/HEAD/tdmpc2_jax/world_model.py --------------------------------------------------------------------------------