├── .gitignore ├── .pre-commit-config.yaml ├── README.md ├── configs ├── base_config.py ├── bridgedata_config.py ├── calvin_config.py ├── real_config.py └── state_config.py ├── jaxrl_m ├── agents │ ├── __init__.py │ └── continuous │ │ ├── action_optimization.py │ │ ├── auto_regressive_transformer.py │ │ ├── base_policy.py │ │ ├── bc.py │ │ ├── calql.py │ │ ├── cql.py │ │ ├── ddpm_bc.py │ │ ├── diffusion_q_learning.py │ │ ├── iql.py │ │ ├── openvla.py │ │ ├── parl_calql.py │ │ ├── parl_iql.py │ │ └── sac.py ├── common │ ├── common.py │ ├── encoding.py │ ├── evaluation.py │ ├── initialization.py │ ├── optimizers.py │ ├── traj.py │ ├── typing.py │ └── wandb.py ├── data │ ├── bridge_dataset.py │ ├── calvin_dataset.py │ ├── calvin_to_lerobot.py │ ├── create_metaworld_dataset.py │ ├── dataset.py │ ├── goal_conditioned_buffer.py │ ├── image_replay_buffer.py │ ├── np_goal_relabeling.py │ ├── replay_buffer.py │ ├── roboverse_dataset.py │ ├── text_processing.py │ ├── tf_augmentations.py │ └── tf_goal_relabeling.py ├── envs │ ├── calvin.py │ ├── calvin_config │ │ ├── calvin.yaml │ │ └── calvin_conf │ │ │ ├── cameras │ │ │ ├── cameras │ │ │ │ ├── gripper.yaml │ │ │ │ ├── opposing.yaml │ │ │ │ ├── static.yaml │ │ │ │ └── tactile.yaml │ │ │ ├── no_cameras.yaml │ │ │ ├── static_and_gripper.yaml │ │ │ ├── static_and_tactile.yaml │ │ │ ├── static_gripper_tactile.yaml │ │ │ └── static_only.yaml │ │ │ ├── config_data_collection.yaml │ │ │ ├── config_rendering.yaml │ │ │ ├── digit_sensor │ │ │ └── config_digit.yml │ │ │ ├── recorder │ │ │ └── recorder.yaml │ │ │ ├── robot │ │ │ ├── panda.yaml │ │ │ ├── panda_digit.yaml │ │ │ └── panda_longer_finger.yaml │ │ │ ├── scene │ │ │ ├── basic_playtable.yaml │ │ │ ├── basic_tabletop.yaml │ │ │ ├── calvin_scene_A.yaml │ │ │ ├── calvin_scene_B.yaml │ │ │ ├── calvin_scene_C.yaml │ │ │ ├── calvin_scene_D.yaml │ │ │ ├── calvin_scene_D_eval.yaml │ │ │ ├── empty_playtable.yaml │ │ │ ├── tabletop_1.yaml │ │ │ ├── tabletop_2.yaml │ │ │ └── tabletop_3.yaml │ │ │ ├── tasks │ │ │ ├── customized_tasks.yaml │ │ │ ├── new_playtable_tasks.yaml │ │ │ └── play_table_tasks.yaml │ │ │ └── vr_input │ │ │ ├── vr_controller │ │ │ ├── oculus.yaml │ │ │ └── vive.yaml │ │ │ └── vr_input.yaml │ ├── d4rl.py │ ├── discretize_continuous_action.py │ ├── franka_env.py │ ├── metaworld.py │ ├── pointmaze.py │ ├── procgen.py │ ├── robot_env.py │ ├── robot_manual_reward_functions.py │ └── wrappers │ │ ├── chunking.py │ │ ├── dmcgym.py │ │ ├── mujoco.py │ │ ├── norm.py │ │ ├── remap.py │ │ ├── reward_scale.py │ │ ├── roboverse.py │ │ └── video_recorder.py ├── networks │ ├── actor_critic_nets.py │ ├── diffusion_nets.py │ ├── discrete_nets.py │ ├── distributional.py │ ├── lagrange.py │ └── mlp.py ├── transformers │ ├── models │ │ ├── __init__.py │ │ ├── infrastructure.py │ │ └── q_transformer.py │ ├── q_functions.py │ └── utils │ │ ├── __init__.py │ │ └── misc.py ├── utils │ ├── jax_utils.py │ ├── sim_utils.py │ ├── timer_utils.py │ ├── train_utils.py │ └── visualization_utils.py └── vision │ ├── __init__.py │ ├── bigvision_common.py │ ├── bigvision_resnetv2.py │ ├── bigvision_utils.py │ ├── cvae.py │ ├── data_augmentations.py │ ├── film_conditioning_layer.py │ ├── impala.py │ ├── resnet_v1.py │ ├── robofume_encoders.py │ └── small_encoders.py ├── requirements-tpu.txt ├── requirements.txt ├── scripts ├── bridgedata_numpy_to_tfrecord.py ├── bridgedata_raw_to_numpy.py ├── maze2d_utils.py └── openvla_caching_worker.py ├── setup.py └── train.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/.gitignore -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/README.md -------------------------------------------------------------------------------- /configs/base_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/configs/base_config.py -------------------------------------------------------------------------------- /configs/bridgedata_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/configs/bridgedata_config.py -------------------------------------------------------------------------------- /configs/calvin_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/configs/calvin_config.py -------------------------------------------------------------------------------- /configs/real_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/configs/real_config.py -------------------------------------------------------------------------------- /configs/state_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/configs/state_config.py -------------------------------------------------------------------------------- /jaxrl_m/agents/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/__init__.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/action_optimization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/action_optimization.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/auto_regressive_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/auto_regressive_transformer.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/base_policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/base_policy.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/bc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/bc.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/calql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/calql.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/cql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/cql.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/ddpm_bc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/ddpm_bc.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/diffusion_q_learning.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/diffusion_q_learning.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/iql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/iql.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/openvla.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/openvla.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/parl_calql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/parl_calql.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/parl_iql.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/parl_iql.py -------------------------------------------------------------------------------- /jaxrl_m/agents/continuous/sac.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/agents/continuous/sac.py -------------------------------------------------------------------------------- /jaxrl_m/common/common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/common.py -------------------------------------------------------------------------------- /jaxrl_m/common/encoding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/encoding.py -------------------------------------------------------------------------------- /jaxrl_m/common/evaluation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/evaluation.py -------------------------------------------------------------------------------- /jaxrl_m/common/initialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/initialization.py -------------------------------------------------------------------------------- /jaxrl_m/common/optimizers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/optimizers.py -------------------------------------------------------------------------------- /jaxrl_m/common/traj.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/traj.py -------------------------------------------------------------------------------- /jaxrl_m/common/typing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/typing.py -------------------------------------------------------------------------------- /jaxrl_m/common/wandb.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/common/wandb.py -------------------------------------------------------------------------------- /jaxrl_m/data/bridge_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/bridge_dataset.py -------------------------------------------------------------------------------- /jaxrl_m/data/calvin_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/calvin_dataset.py -------------------------------------------------------------------------------- /jaxrl_m/data/calvin_to_lerobot.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/calvin_to_lerobot.py -------------------------------------------------------------------------------- /jaxrl_m/data/create_metaworld_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/create_metaworld_dataset.py -------------------------------------------------------------------------------- /jaxrl_m/data/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/dataset.py -------------------------------------------------------------------------------- /jaxrl_m/data/goal_conditioned_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/goal_conditioned_buffer.py -------------------------------------------------------------------------------- /jaxrl_m/data/image_replay_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/image_replay_buffer.py -------------------------------------------------------------------------------- /jaxrl_m/data/np_goal_relabeling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/np_goal_relabeling.py -------------------------------------------------------------------------------- /jaxrl_m/data/replay_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/replay_buffer.py -------------------------------------------------------------------------------- /jaxrl_m/data/roboverse_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/roboverse_dataset.py -------------------------------------------------------------------------------- /jaxrl_m/data/text_processing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/text_processing.py -------------------------------------------------------------------------------- /jaxrl_m/data/tf_augmentations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/tf_augmentations.py -------------------------------------------------------------------------------- /jaxrl_m/data/tf_goal_relabeling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/data/tf_goal_relabeling.py -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin.py -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/gripper.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/gripper.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/opposing.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/opposing.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/static.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/static.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/tactile.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/cameras/tactile.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/no_cameras.yaml: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_and_gripper.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_and_gripper.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_and_tactile.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_and_tactile.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_gripper_tactile.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_gripper_tactile.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_only.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/cameras/static_only.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/config_data_collection.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/config_data_collection.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/config_rendering.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/config_rendering.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/digit_sensor/config_digit.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/digit_sensor/config_digit.yml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/recorder/recorder.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/recorder/recorder.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/robot/panda.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/robot/panda.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/robot/panda_digit.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/robot/panda_digit.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/robot/panda_longer_finger.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/robot/panda_longer_finger.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/basic_playtable.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/basic_playtable.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/basic_tabletop.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/basic_tabletop.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_A.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_A.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_B.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_B.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_C.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_C.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_D.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_D.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_D_eval.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/calvin_scene_D_eval.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/empty_playtable.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/empty_playtable.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_1.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_1.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_2.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_2.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_3.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/scene/tabletop_3.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/tasks/customized_tasks.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/tasks/customized_tasks.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/tasks/new_playtable_tasks.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/tasks/new_playtable_tasks.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/tasks/play_table_tasks.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/tasks/play_table_tasks.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_controller/oculus.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_controller/oculus.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_controller/vive.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_controller/vive.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_input.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/calvin_config/calvin_conf/vr_input/vr_input.yaml -------------------------------------------------------------------------------- /jaxrl_m/envs/d4rl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/d4rl.py -------------------------------------------------------------------------------- /jaxrl_m/envs/discretize_continuous_action.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/discretize_continuous_action.py -------------------------------------------------------------------------------- /jaxrl_m/envs/franka_env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/franka_env.py -------------------------------------------------------------------------------- /jaxrl_m/envs/metaworld.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/metaworld.py -------------------------------------------------------------------------------- /jaxrl_m/envs/pointmaze.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/pointmaze.py -------------------------------------------------------------------------------- /jaxrl_m/envs/procgen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/procgen.py -------------------------------------------------------------------------------- /jaxrl_m/envs/robot_env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/robot_env.py -------------------------------------------------------------------------------- /jaxrl_m/envs/robot_manual_reward_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/robot_manual_reward_functions.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/chunking.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/chunking.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/dmcgym.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/dmcgym.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/mujoco.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/mujoco.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/norm.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/remap.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/remap.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/reward_scale.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/reward_scale.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/roboverse.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/roboverse.py -------------------------------------------------------------------------------- /jaxrl_m/envs/wrappers/video_recorder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/envs/wrappers/video_recorder.py -------------------------------------------------------------------------------- /jaxrl_m/networks/actor_critic_nets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/actor_critic_nets.py -------------------------------------------------------------------------------- /jaxrl_m/networks/diffusion_nets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/diffusion_nets.py -------------------------------------------------------------------------------- /jaxrl_m/networks/discrete_nets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/discrete_nets.py -------------------------------------------------------------------------------- /jaxrl_m/networks/distributional.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/distributional.py -------------------------------------------------------------------------------- /jaxrl_m/networks/lagrange.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/lagrange.py -------------------------------------------------------------------------------- /jaxrl_m/networks/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/networks/mlp.py -------------------------------------------------------------------------------- /jaxrl_m/transformers/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /jaxrl_m/transformers/models/infrastructure.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/transformers/models/infrastructure.py -------------------------------------------------------------------------------- /jaxrl_m/transformers/models/q_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/transformers/models/q_transformer.py -------------------------------------------------------------------------------- /jaxrl_m/transformers/q_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/transformers/q_functions.py -------------------------------------------------------------------------------- /jaxrl_m/transformers/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .misc import * -------------------------------------------------------------------------------- /jaxrl_m/transformers/utils/misc.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/transformers/utils/misc.py -------------------------------------------------------------------------------- /jaxrl_m/utils/jax_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/utils/jax_utils.py -------------------------------------------------------------------------------- /jaxrl_m/utils/sim_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/utils/sim_utils.py -------------------------------------------------------------------------------- /jaxrl_m/utils/timer_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/utils/timer_utils.py -------------------------------------------------------------------------------- /jaxrl_m/utils/train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/utils/train_utils.py -------------------------------------------------------------------------------- /jaxrl_m/utils/visualization_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/utils/visualization_utils.py -------------------------------------------------------------------------------- /jaxrl_m/vision/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/__init__.py -------------------------------------------------------------------------------- /jaxrl_m/vision/bigvision_common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/bigvision_common.py -------------------------------------------------------------------------------- /jaxrl_m/vision/bigvision_resnetv2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/bigvision_resnetv2.py -------------------------------------------------------------------------------- /jaxrl_m/vision/bigvision_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/bigvision_utils.py -------------------------------------------------------------------------------- /jaxrl_m/vision/cvae.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/cvae.py -------------------------------------------------------------------------------- /jaxrl_m/vision/data_augmentations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/data_augmentations.py -------------------------------------------------------------------------------- /jaxrl_m/vision/film_conditioning_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/film_conditioning_layer.py -------------------------------------------------------------------------------- /jaxrl_m/vision/impala.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/impala.py -------------------------------------------------------------------------------- /jaxrl_m/vision/resnet_v1.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/resnet_v1.py -------------------------------------------------------------------------------- /jaxrl_m/vision/robofume_encoders.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/robofume_encoders.py -------------------------------------------------------------------------------- /jaxrl_m/vision/small_encoders.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/jaxrl_m/vision/small_encoders.py -------------------------------------------------------------------------------- /requirements-tpu.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/requirements-tpu.txt -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/requirements.txt -------------------------------------------------------------------------------- /scripts/bridgedata_numpy_to_tfrecord.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/scripts/bridgedata_numpy_to_tfrecord.py -------------------------------------------------------------------------------- /scripts/bridgedata_raw_to_numpy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/scripts/bridgedata_raw_to_numpy.py -------------------------------------------------------------------------------- /scripts/maze2d_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/scripts/maze2d_utils.py -------------------------------------------------------------------------------- /scripts/openvla_caching_worker.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/scripts/openvla_caching_worker.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/setup.py -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MaxSobolMark/PolicyAgnosticRL/HEAD/train.py --------------------------------------------------------------------------------