├── .gitignore ├── LICENSE ├── README.md ├── contributing.md ├── examples ├── mujoco │ ├── d4rl_utils.py │ ├── iql.py │ ├── run_d4rl_iql.py │ ├── run_mujoco_sac.py │ └── sac.py └── vision_agents │ ├── continuous_bc.py │ ├── discrete_bc.py │ └── discrete_cql.py ├── jaxrl_m ├── common.py ├── dataset.py ├── evaluation.py ├── networks.py ├── typing.py ├── vision │ ├── README.md │ ├── __init__.py │ ├── bridge_resnet_v1.py │ ├── data_augmentations.py │ ├── impala.py │ ├── mae.py │ ├── preprocess.py │ ├── pretrained_utils.py │ ├── resnet_v1.py │ ├── resnet_v2.py │ ├── small_encoders.py │ ├── vit.py │ └── vit_resnet.py └── wandb.py └── setup.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/README.md -------------------------------------------------------------------------------- /contributing.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/contributing.md -------------------------------------------------------------------------------- /examples/mujoco/d4rl_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/mujoco/d4rl_utils.py -------------------------------------------------------------------------------- /examples/mujoco/iql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/mujoco/iql.py -------------------------------------------------------------------------------- /examples/mujoco/run_d4rl_iql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/mujoco/run_d4rl_iql.py -------------------------------------------------------------------------------- /examples/mujoco/run_mujoco_sac.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/mujoco/run_mujoco_sac.py -------------------------------------------------------------------------------- /examples/mujoco/sac.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/mujoco/sac.py -------------------------------------------------------------------------------- /examples/vision_agents/continuous_bc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/vision_agents/continuous_bc.py -------------------------------------------------------------------------------- /examples/vision_agents/discrete_bc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/vision_agents/discrete_bc.py -------------------------------------------------------------------------------- /examples/vision_agents/discrete_cql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/examples/vision_agents/discrete_cql.py -------------------------------------------------------------------------------- /jaxrl_m/common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/common.py -------------------------------------------------------------------------------- /jaxrl_m/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/dataset.py -------------------------------------------------------------------------------- /jaxrl_m/evaluation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/evaluation.py -------------------------------------------------------------------------------- /jaxrl_m/networks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/networks.py -------------------------------------------------------------------------------- /jaxrl_m/typing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/typing.py -------------------------------------------------------------------------------- /jaxrl_m/vision/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/README.md -------------------------------------------------------------------------------- /jaxrl_m/vision/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/__init__.py -------------------------------------------------------------------------------- /jaxrl_m/vision/bridge_resnet_v1.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/bridge_resnet_v1.py -------------------------------------------------------------------------------- /jaxrl_m/vision/data_augmentations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/data_augmentations.py -------------------------------------------------------------------------------- /jaxrl_m/vision/impala.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/impala.py -------------------------------------------------------------------------------- /jaxrl_m/vision/mae.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/mae.py -------------------------------------------------------------------------------- /jaxrl_m/vision/preprocess.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/preprocess.py -------------------------------------------------------------------------------- /jaxrl_m/vision/pretrained_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/pretrained_utils.py -------------------------------------------------------------------------------- /jaxrl_m/vision/resnet_v1.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/resnet_v1.py -------------------------------------------------------------------------------- /jaxrl_m/vision/resnet_v2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/resnet_v2.py -------------------------------------------------------------------------------- /jaxrl_m/vision/small_encoders.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/small_encoders.py -------------------------------------------------------------------------------- /jaxrl_m/vision/vit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/vit.py -------------------------------------------------------------------------------- /jaxrl_m/vision/vit_resnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/vision/vit_resnet.py -------------------------------------------------------------------------------- /jaxrl_m/wandb.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/jaxrl_m/wandb.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dibyaghosh/jaxrl_m/HEAD/setup.py --------------------------------------------------------------------------------