# Overview
OGBench is a benchmark designed to facilitate algorithms research in offline goal-conditioned reinforcement learning (RL),
offline unsupervised RL, and offline RL.
See the [project page](https://seohong.me/projects/ogbench/) for videos and more details about the environments, tasks, and datasets.
### Features
- **8 types** of realistic and diverse environments ([videos](https://seohong.me/projects/ogbench/)):
- **Locomotion**: PointMaze, AntMaze, HumanoidMaze, and AntSoccer.
- **Manipulation**: Cube, Scene, and Puzzle.
- **Drawing**: Powderworld.
- **85 datasets** covering various challenges in offline goal-conditioned RL.
- **410 tasks** for standard (i.e., non-goal-conditioned) offline RL.
- Support for both **pixel-based** and **state-based** observations.
- **Clean, well-tuned reference implementations** of 6 offline goal-conditioned RL algorithms
(GCBC, GCIVL, GCIQL, QRL, CRL, and HIQL) based on JAX.
- **Fully reproducible** scripts for [the entire benchmark table](impls/hyperparameters.sh)
and [datasets](data_gen_scripts/commands.sh).
- `pip`-installable, easy-to-use APIs based on Gymnasium.
- No major dependencies other than MuJoCo.
# Quick Start
### Installation
OGBench can be easily installed via PyPI:
```shell
pip install ogbench
```
It requires Python 3.8+ and has only three dependencies: `mujoco >= 3.1.6`, `dm_control >= 1.0.20`,
and `gymnasium`.
To use OGBench for **offline goal-conditioned RL**,
go to [this section](#usage-for-offline-goal-conditioned-rl).
To use OGBench for **standard (non-goal-conditioned) offline RL**,
go to [this section](#usage-for-standard-non-goal-conditioned-offline-rl).
### Usage for offline goal-conditioned RL
After installing OGBench, you can create an environment and datasets using `ogbench.make_env_and_datasets`.
The environment follows the [Gymnasium](https://gymnasium.farama.org/) interface.
The datasets will be automatically downloaded during the first run.
Here is an example of how to use OGBench for offline goal-conditioned RL:
> [!CAUTION]
> Do **not** use `gymnasium.make` to create an environment. Use `ogbench.make_env_and_datasets` instead.
> To create an environment without loading datasets, use `env_only=True` in `ogbench.make_env_and_datasets`.
```python
import ogbench
# Make an environment and datasets (they will be automatically downloaded).
dataset_name = 'humanoidmaze-large-navigate-v0'
env, train_dataset, val_dataset = ogbench.make_env_and_datasets(dataset_name)
# Train your offline goal-conditioned RL agent on the dataset.
# ...
# Evaluate the agent.
for task_id in [1, 2, 3, 4, 5]:
# Reset the environment and set the evaluation task.
ob, info = env.reset(
options=dict(
task_id=task_id, # Set the evaluation task. Each environment provides five
# evaluation goals, and `task_id` must be in [1, 5].
render_goal=True, # Set to `True` to get a rendered goal image (optional).
)
)
goal = info['goal'] # Get the goal observation to pass to the agent.
goal_rendered = info['goal_rendered'] # Get the rendered goal image (optional).
done = False
while not done:
action = env.action_space.sample() # Replace this with your agent's action.
ob, reward, terminated, truncated, info = env.step(action) # Gymnasium-style step.
# If the agent reaches the goal, `terminated` will be `True`. If the episode length
# exceeds the maximum length without reaching the goal, `truncated` will be `True`.
# `reward` is 1 if the agent reaches the goal and 0 otherwise.
done = terminated or truncated
frame = env.render() # Render the current frame (optional).
success = info['success'] # Whether the agent reached the goal (0 or 1).
# `terminated` also indicates this.
```
You can find a complete example of a training script for offline goal-conditioned RL in the `impls` directory.
See the next section for more details on the reference implementations.
### Usage for standard (non-goal-conditioned) offline RL
OGBench also provides single-task variants of the environments for standard (reward-maximizing) offline RL.
Each locomotion and manipulation environment provides five different single-task tasks corresponding to the five evaluation goals,
and they are named with the suffix `singletask-task[n]` (e.g., `scene-play-singletask-task2-v0`),
where `[n]` denotes a number between 1 and 5 (inclusive).
Among the five tasks in each environment,
the most representative one is chosen as the "default" task,
and is *aliased* by the suffix `singletask` without a task number.
Default tasks can be useful for reducing the number of benchmarking environments
or for tuning hyperparameters.
Click to see the list of default tasks
| Environment | Default Task |
|:-------------------:|:------------:|
| `pointmaze-*` | `task1` |
| `antmaze-*` | `task1` |
| `humanoidmaze-*` | `task1` |
| `antsoccer-*` | `task4` |
| `cube-*` | `task2` |
| `scene-*` | `task2` |
| `puzzle-{3x3, 4x4}` | `task4` |
| `puzzle-{4x5, 4x6}` | `task2` |
Here is an example of how to use OGBench for standard (non-goal-conditioned) offline RL:
> [!CAUTION]
> Do **not** use `gymnasium.make` to create an environment. Use `ogbench.make_env_and_datasets` instead.
> To create an environment without loading datasets, use `env_only=True` in `ogbench.make_env_and_datasets`.
> [!NOTE]
> Offline RL datasets contain both the `terminals` and `masks` fields.
>
> * `masks` denotes whether the agent should get a Bellman backup from the next observation.
> It is 0 only when the task is complete (and 1 otherwise).
> In this case, the agent should set the target Q-value to 0,
> instead of using the next observation's target Q-value.
> * `terminals` simply denotes whether the dataset trajectory is over,
> regardless of task completion.
>
> For example, in `antmaze-large-navigate-singletask-v0`, the dataset contains 1M transitions,
> with each trajectory having a length of 1000.
> Hence, `sum(dataset['terminals'])` is exactly 1000 (i.e., 1 at the end of each trajectory),
> whereas `sum(dataset['masks'])` can vary
> depending on how many times the agent reaches the goal.
> Note that dataset trajectories do not terminate even when the agent reaches the goal,
> as they are collected by a scripted policy that is not task-aware.
>
> For standard Q-learning, you likely only need `masks`,
> but for other trajectory-aware algorithms (e.g., hierarchical RL or trajectory modeling-based approaches),
> you may need both `masks` and `terminals`.
> See [the IQL implementation in the FQL repository](https://github.com/seohongpark/fql/blob/master/agents/iql.py)
> for an example of how to use `masks`.
```python
import ogbench
# Make an environment and datasets (they will be automatically downloaded).
# In `cube-double`, the default task is `task2`, and it is also callable by
# `cube-double-play-singletask-v0`.
dataset_name = 'cube-double-play-singletask-task2-v0'
env, train_dataset, val_dataset = ogbench.make_env_and_datasets(dataset_name)
# Train your offline RL agent on the dataset.
# ...
# Evaluate the agent.
ob, info = env.reset() # Reset the environment.
done = False
while not done:
action = env.action_space.sample() # Replace this with your agent's action.
ob, reward, terminated, truncated, info = env.step(action) # Gymnasium-style step.
# If the agent achieves the task, `terminated` will be `True`. If the episode length
# exceeds the maximum length without achieving the task, `truncated` will be `True`.
done = terminated or truncated
frame = env.render() # Render the current frame (optional).
success = info['success'] # Whether the agent achieved the task (0 or 1).
```
For standard offline RL, we do not provide official reference implementations or benchmarking results.
However, you may find implementations of some offline RL algorithms (e.g., IQL, ReBRAC, and FQL) with partial benchmarking results
in [this repository](https://github.com/seohongpark/fql).
### Dataset APIs
OGBench provides several APIs to download and load datasets.
The simplest way is to use `ogbench.make_env_and_datasets` as shown above,
which creates an environment and loads training and validation datasets.
The datasets will automatically be downloaded to the directory specified by `dataset_dir` during the first run
(default: `~/.ogbench/data`).
`ogbench.make_env_and_datasets` also provides the `compact_dataset` option,
which returns a dataset without the `next_observations` field (see below).
For example:
```python
import ogbench
# Make an environment and load datasets.
dataset_name = 'antmaze-large-navigate-v0'
env, train_dataset, val_dataset = ogbench.make_env_and_datasets(
dataset_name, # Dataset name.
dataset_dir='~/.ogbench/data', # Directory to save datasets (optional).
compact_dataset=False, # Whether to use a compact dataset (optional; see below).
)
# Assume each dataset trajectory has a length of 4, and (s0, a0, s1), (s1, a1, s2),
# (s2, a2, s3), (s3, a3, s4) are the transition tuples.
# If `compact_dataset` is `False`, the dataset will have the following structure:
# |<- traj 1 ->| |<- traj 2 ->| ...
# ----------------------------------------------------------
# 'observations' : [s0, s1, s2, s3, s0, s1, s2, s3, ...]
# 'actions' : [a0, a1, a2, a3, a0, a1, a2, a3, ...]
# 'next_observations': [s1, s2, s3, s4, s1, s2, s3, s4, ...]
# 'terminals' : [ 0, 0, 0, 1, 0, 0, 0, 1, ...]
# If `compact_dataset` is `True`, the dataset will have the following structure, where the
# `next_observations` field is omitted. Instead, it includes a `valids` field indicating
# whether the next observation is valid:
# |<--- traj 1 --->| |<--- traj 2 --->| ...
# ------------------------------------------------------------------
# 'observations' : [s0, s1, s2, s3, s4, s0, s1, s2, s3, s4, ...]
# 'actions' : [a0, a1, a2, a3, a4, a0, a1, a2, a3, a4, ...]
# 'terminals' : [ 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, ...]
# 'valids' : [ 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, ...]
```
To download multiple datasets at once, you can use `ogbench.download_datasets`:
```python
import ogbench
dataset_names = [
'humanoidmaze-medium-navigate-v0',
'visual-puzzle-3x3-play-v0',
'powderworld-easy-play-v0',
]
ogbench.download_datasets(
dataset_names, # List of dataset names.
dataset_dir='~/.ogbench/data', # Directory to save datasets (optional).
)
```
# Reference Implementations
OGBench also provides JAX-based reference implementations of six offline goal-conditioned RL algorithms
(GCBC, GCIVL, GCIQL, QRL, CRL and HIQL).
They are provided in the `impls` directory as a **standalone** codebase.
You can safely remove the other parts of the repository if you only need the reference implementations
and do not want to modify the environments.
### Installation
Our reference implementations require Python 3.9+ and additional dependencies, including `jax >= 0.4.26`.
To install these dependencies, run:
```shell
cd impls
pip install -r requirements.txt
```
By default, it uses the PyPI version of OGBench.
If you want to use a local version of OGBench (e.g., for training methods on modified environments),
run instead `pip install -e ".[train]"` in the root directory.
### Running the reference implementations
Each algorithm is implemented in a separate file in the `agents` directory.
We provide implementations of the following offline goal-conditioned RL algorithms:
- `gcbc.py`: Goal-Conditioned Behavioral Cloning (GCBC)
- `gcivl.py`: Goal-Conditioned Implicit V-Learning (GCIVL)
- `gciql.py`: Goal-Conditioned Implicit Q-Learning (GCIQL)
- `qrl.py`: Quasimetric Reinforcement Learning (QRL)
- `crl.py`: Contrastive Reinforcement Learning (CRL)
- `hiql.py`: Hierarchical Implicit Q-Learning (HIQL)
To train an agent, you can run the `main.py` script.
Training metrics, evaluation metrics, and videos are logged via `wandb` by default.
Here are some example commands (see [hyperparameters.sh](impls/hyperparameters.sh) for the full list of commands):
```shell
# antmaze-large-navigate-v0 (GCBC)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/gcbc.py
# antmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/gcivl.py --agent.alpha=10.0
# antmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/gciql.py --agent.alpha=0.3
# antmaze-large-navigate-v0 (QRL)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/qrl.py --agent.alpha=0.003
# antmaze-large-navigate-v0 (CRL)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/crl.py --agent.alpha=0.1
# antmaze-large-navigate-v0 (HIQL)
python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
```
Each run typically takes 2-5 hours (on state-based tasks)
or 5-12 hours (on pixel-based tasks) on a single A5000 GPU.
For large pixel-based datasets (e.g., `visual-puzzle-4x6-play-v0` with 5M transitions),
up to 120GB of RAM may be required.
> [!NOTE]
> If you are running on a remote/headless server without a display, you can use EGL for rendering by setting the `MUJOCO_GL` environment variable:
> ```shell
> MUJOCO_GL=egl python main.py --env_name=antmaze-large-navigate-v0 --agent=agents/gcbc.py
> ```
### Tips for hyperparameters and flags
To reproduce the results in the paper, you need to use the hyperparameters provided.
We provide a complete list of the exact command-line flags used to produce the main benchmark table
in the paper in [hyperparameters.sh](impls/hyperparameters.sh).
Below, we highlight some important hyperparameters and common pitfalls:
- Regardless of the algorithms, one of the most important hyperparameters is `agent.alpha` (i.e., the temperature (AWR) or the BC coefficient (DDPG+BC))
for the actor loss. It is crucial to tune this hyperparameter when running an algorithm on a new environment.
In the paper, we provide a separate table of the policy extraction hyperparameters,
which are individually tuned for each environment and dataset category.
- By default, actor goals are uniformly sampled from the future states in the same trajectory.
We found this works best in most cases, but you can adjust this to allow random actor goals
(e.g., by setting `--agent.actor_p_trajgoal=0.5 --agent.actor_p_randomgoal=0.5`).
This is especially important for datasets that require stitching.
See the hyperparameter table in the paper for the values used in benchmarking.
- For GCIQL, CRL, and QRL, we provide two policy extraction methods: AWR and DDPG+BC.
In general, DDPG+BC works better than AWR (see [this paper](https://arxiv.org/abs/2406.09329) for the reasons),
but DDPG+BC is usually more sensitive to the `alpha` hyperparameter than AWR.
As such, in a new environment, we recommend starting with AWR to get a sence of the performance
and then switching to DDPG+BC to further improve the performance.
- Our QRL implementation provides two quasimetric parameterizations: MRN and IQE.
We found that IQE (default) works better in general, but it is almost twice as slow as MRN.
- In CRL, we found that using `--agent.actor_log_q=True` (which is set by default) is important for strong performance, especially in locomotion environments.
We found this doesn't help much with other algorithms.
- In HIQL, setting `--agent.low_actor_rep_grad=True` (which is `False` by default) is crucial in pixel-based environments.
This allows gradients to flow from the low-level actor loss to the subgoal representation, which helps maintain better representations.
- In pixel-based environments, don't forget to set `agent.encoder`. We used `--agent.encoder=impala_small` across all pixel-based environments.
- In discrete-action environments (e.g., Powderworld), don't forget to set `--agent.discrete=True`.
- In Powderworld, use `--eval_temperature=0.3`, which helps prevent the agent from getting stuck in certain states.
# Reproducing Datasets
We provide the full scripts and exact command-line flags used to produce all the datasets in OGBench.
The scripts are provided in the `data_gen_scripts` directory.
### Installation
Data-generation scripts for locomotion environments require Python 3.9+ and additional dependencies,
including `jax >= 0.4.26`, to train and load expert agents.
For manipulation and drawing environments, no additional dependencies are required.
To install the necessary dependencies for locomotion environments, run the following command in the root directory:
```shell
pip install -e ".[train]"
```
This installs the same dependencies as the reference implementations, but in the editable mode (`-e`).
### Reproducing datasets
To reproduce datasets, you can run the scripts in the `data_gen_scripts` directory.
For locomotion environments, you need to first download the expert policies.
We provide the exact command-line flags used to produce the datasets in [commands.sh](data_gen_scripts/commands.sh).
Here is an example of how to reproduce a dataset for the `antmaze-large-navigate-v0` task:
```shell
cd data_gen_scripts
# Download the expert policies for locomotion environments (not required for other environments).
wget https://rail.eecs.berkeley.edu/datasets/ogbench/experts.tar.gz
tar xf experts.tar.gz && rm experts.tar.gz
# Create a directory to save datasets.
mkdir -p data
# Add the `impls` directory to PYTHONPATH.
# Alternatively, you can move the contents of `data_gen_scripts` to `impls` instead of setting PYTHONPATH.
export PYTHONPATH="../impls:${PYTHONPATH}"
# Generate a dataset for `antmaze-large-navigate-v0`.
python generate_locomaze.py --env_name=antmaze-large-v0 --save_path=data/antmaze-large-navigate-v0.npz
```
### Reproducing expert policies
If you want to train your own expert policies from scratch, you can run the corresponding commands in [commands.sh](data_gen_scripts/commands.sh).
For example, to train an Ant expert policy, you can run the following command in the `data_gen_scripts` directory after setting `PYTHONPATH` as above:
```shell
python main_sac.py --env_name=online-ant-xy-v0
```
# Additional Features
- We support `-oraclerep` variants, which provide ground-truth goal representations
(e.g., in `antmaze-large-navigate-oraclerep-v0`,
the goal is defined only by the x-y position, not including the agent's proprioceptive states).
- We also provide the `cube-octuple` task, which involves eight cubes.
While we do not provide a default dataset for this task, you may download the 100M-sized dataset below.
- For some tasks, we provide larger datasets with 100M transitions, collected by the same scripted policy as the original datasets.
They can be manually downloaded from the following links (see [this repository](https://github.com/seohongpark/horizon-reduction) for examples of how to load these datasets):
- `humanoidmaze-giant-navigate-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/humanoidmaze-giant-navigate-100m-v0
- `cube-double-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/cube-double-play-100m-v0
- `cube-triple-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/cube-triple-play-100m-v0
- `cube-quadruple-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/cube-quadruple-play-100m-v0
- `cube-quadruple-noisy-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/cube-quadruple-noisy-100m-v0
- `cube-octuple-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/cube-octuple-play-100m-v0
- `scene-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/scene-play-100m-v0
- `puzzle-3x3-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/puzzle-3x3-play-100m-v0
- `puzzle-4x4-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/puzzle-4x4-play-100m-v0
- `puzzle-4x5-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/puzzle-4x5-play-100m-v0
- `puzzle-4x6-play-100m-v0`: https://rail.eecs.berkeley.edu/datasets/ogbench/puzzle-4x6-play-100m-v0
# Caveats
- Starting from OGBench 1.2.0, `singletask` environments compute `reward`, `terminated`, and `info['success']`
based on the current state (i.e., compute `r(s)` instead of `r(s')` for an `(s, a, s')` tuple)
to be consistent with the dataset reward structure.
In earlier versions, they were computed based on the next state (`s'`),
so this change may lead to slight differences in evaluation results (though we expect the differences to be negligible).
You can set `success_timing='post'` in `ogbench.make_env_and_datasets` to restore the previous behavior if needed.
We also note that this change only affects `singletask` environments; goal-conditioned environments remain unchanged
(they always compute `terminated` and `info['success']` based on `s'` even in the latest version).
# Acknowledgments
This codebase is inspired by or partly uses code from the following repositories:
- [D4RL](https://github.com/Farama-Foundation/D4RL) for the dataset structure and the AntMaze environment.
- [Gymnasium](https://github.com/Farama-Foundation/Gymnasium) and [dm_control](https://github.com/google-deepmind/dm_control) for the agents (Ant and Humanoid) in the locomotion environments.
- [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie) for the robot descriptions (Universal Robots UR5e and Robotiq 2F-85) in the manipulation environments.
- [jaxlie](https://github.com/brentyi/jaxlie) for Lie group operations in the manipulation environments.
- [Meta-World](https://github.com/Farama-Foundation/Metaworld) for the objects (drawer, window, and button) in the manipulation environments.
- [Powderworld](https://github.com/kvfrans/powderworld) for the Powderworld environment.
- [NumPyConv2D](https://github.com/99991/NumPyConv2D) for the NumPy Conv2D implementation in the Powderworld environment.
- [jaxrl_m](https://github.com/dibyaghosh/jaxrl_m), [rlbase](https://github.com/kvfrans/rlbase_stable),
[HIQL](https://github.com/seohongpark/HIQL), and [cmd-notebook](https://github.com/vivekmyers/cmd-notebook)
for JAX-based implementations of RL algorithms.
Special thanks to [Kevin Zakka](https://kzakka.com/) for providing the initial codebase for the manipulation environments.
# Citation
```bibtex
@inproceedings{ogbench_park2025,
title={OGBench: Benchmarking Offline Goal-Conditioned RL},
author={Park, Seohong and Frans, Kevin and Eysenbach, Benjamin and Levine, Sergey},
booktitle={International Conference on Learning Representations (ICLR)},
year={2025},
}
```
================================================
FILE: data_gen_scripts/commands.sh
================================================
# Commands to train expert policies.
# ant (online-ant-xy-v0)
python main_sac.py --env_name=online-ant-xy-v0 --train_steps=400000 --eval_interval=100000 --save_interval=400000 --log_interval=5000
# antball (online-antball-v0)
python main_sac.py --env_name=online-antball-v0 --train_steps=12000000 --train_interval=4 --eval_interval=500000 --save_interval=12000000 --log_interval=20000 --agent.layer_norm=True --terminate_at_end=1
# humanoid (online-humanoid-xy-v0)
python main_sac.py --env_name=online-humanoid-xy-v0 --train_steps=40000000 --train_interval=4 --eval_interval=500000 --save_interval=40000000 --log_interval=20000 --agent.value_hidden_dims="(1024, 1024, 1024)" --agent.layer_norm=True --agent.min_q=False
# Commands to reproduce datasets.
# pointmaze-medium-navigate-v0
python generate_locomaze.py --env_name=pointmaze-medium-v0 --save_path=data/pointmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --noise=0.5
# pointmaze-large-navigate-v0
python generate_locomaze.py --env_name=pointmaze-large-v0 --save_path=data/pointmaze-large-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --noise=0.5
# pointmaze-giant-navigate-v0
python generate_locomaze.py --env_name=pointmaze-giant-v0 --save_path=data/pointmaze-giant-navigate-v0.npz --dataset_type=navigate --num_episodes=500 --max_episode_steps=2001 --noise=0.5
# pointmaze-teleport-navigate-v0
python generate_locomaze.py --env_name=pointmaze-teleport-v0 --save_path=data/pointmaze-teleport-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --noise=0.5
# pointmaze-medium-stitch-v0
python generate_locomaze.py --env_name=pointmaze-medium-v0 --save_path=data/pointmaze-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
# pointmaze-large-stitch-v0
python generate_locomaze.py --env_name=pointmaze-large-v0 --save_path=data/pointmaze-large-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
# pointmaze-giant-stitch-v0
python generate_locomaze.py --env_name=pointmaze-giant-v0 --save_path=data/pointmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
# pointmaze-teleport-stitch-v0
python generate_locomaze.py --env_name=pointmaze-teleport-v0 --save_path=data/pointmaze-teleport-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
# antmaze-medium-navigate-v0
python generate_locomaze.py --env_name=antmaze-medium-v0 --save_path=data/antmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# antmaze-large-navigate-v0
python generate_locomaze.py --env_name=antmaze-large-v0 --save_path=data/antmaze-large-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# antmaze-giant-navigate-v0
python generate_locomaze.py --env_name=antmaze-giant-v0 --save_path=data/antmaze-giant-navigate-v0.npz --dataset_type=navigate --num_episodes=500 --max_episode_steps=2001 --restore_path=experts/ant --restore_epoch=400000
# antmaze-teleport-navigate-v0
python generate_locomaze.py --env_name=antmaze-teleport-v0 --save_path=data/antmaze-teleport-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# antmaze-medium-stitch-v0
python generate_locomaze.py --env_name=antmaze-medium-v0 --save_path=data/antmaze-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# antmaze-large-stitch-v0
python generate_locomaze.py --env_name=antmaze-large-v0 --save_path=data/antmaze-large-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# antmaze-giant-stitch-v0
python generate_locomaze.py --env_name=antmaze-giant-v0 --save_path=data/antmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# antmaze-teleport-stitch-v0
python generate_locomaze.py --env_name=antmaze-teleport-v0 --save_path=data/antmaze-teleport-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# antmaze-medium-explore-v0
python generate_locomaze.py --env_name=antmaze-medium-v0 --save_path=data/antmaze-medium-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# antmaze-large-explore-v0
python generate_locomaze.py --env_name=antmaze-large-v0 --save_path=data/antmaze-large-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# antmaze-teleport-explore-v0
python generate_locomaze.py --env_name=antmaze-teleport-v0 --save_path=data/antmaze-teleport-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# humanoidmaze-medium-navigate-v0
python generate_locomaze.py --env_name=humanoidmaze-medium-v0 --save_path=data/humanoidmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=2001 --restore_path=experts/humanoid --restore_epoch=40000000
# humanoidmaze-large-navigate-v0
python generate_locomaze.py --env_name=humanoidmaze-large-v0 --save_path=data/humanoidmaze-large-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=2001 --restore_path=experts/humanoid --restore_epoch=40000000
# humanoidmaze-giant-navigate-v0
python generate_locomaze.py --env_name=humanoidmaze-giant-v0 --save_path=data/humanoidmaze-giant-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=4001 --restore_path=experts/humanoid --restore_epoch=40000000
# humanoidmaze-medium-stitch-v0
python generate_locomaze.py --env_name=humanoidmaze-medium-v0 --save_path=data/humanoidmaze-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# humanoidmaze-large-stitch-v0
python generate_locomaze.py --env_name=humanoidmaze-large-v0 --save_path=data/humanoidmaze-large-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# humanoidmaze-giant-stitch-v0
python generate_locomaze.py --env_name=humanoidmaze-giant-v0 --save_path=data/humanoidmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=10000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# antsoccer-arena-navigate-v0
python generate_antsoccer.py --env_name=antsoccer-arena-v0 --save_path=data/antsoccer-arena-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --loco_restore_path=experts/ant --loco_restore_epoch=400000 --ball_restore_path=experts/antball --ball_restore_epoch=12000000
# antsoccer-medium-navigate-v0
python generate_antsoccer.py --env_name=antsoccer-medium-v0 --save_path=data/antsoccer-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=4000 --max_episode_steps=1001 --loco_restore_path=experts/ant --loco_restore_epoch=400000 --ball_restore_path=experts/antball --ball_restore_epoch=12000000
# antsoccer-arena-stitch-v0
python generate_antsoccer.py --env_name=antsoccer-arena-v0 --save_path=data/antsoccer-arena-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --loco_restore_path=experts/ant --loco_restore_epoch=400000 --ball_restore_path=experts/antball --ball_restore_epoch=12000000
# antsoccer-medium-stitch-v0
python generate_antsoccer.py --env_name=antsoccer-medium-v0 --save_path=data/antsoccer-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=8000 --max_episode_steps=501 --loco_restore_path=experts/ant --loco_restore_epoch=400000 --ball_restore_path=experts/antball --ball_restore_epoch=12000000
# visual-antmaze-medium-navigate-v0
python generate_locomaze.py --env_name=visual-antmaze-medium-v0 --save_path=data/visual-antmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-large-navigate-v0
python generate_locomaze.py --env_name=visual-antmaze-large-v0 --save_path=data/visual-antmaze-large-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-giant-navigate-v0
python generate_locomaze.py --env_name=visual-antmaze-giant-v0 --save_path=data/visual-antmaze-giant-navigate-v0.npz --dataset_type=navigate --num_episodes=500 --max_episode_steps=2001 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-teleport-navigate-v0
python generate_locomaze.py --env_name=visual-antmaze-teleport-v0 --save_path=data/visual-antmaze-teleport-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-medium-stitch-v0
python generate_locomaze.py --env_name=visual-antmaze-medium-v0 --save_path=data/visual-antmaze-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-large-stitch-v0
python generate_locomaze.py --env_name=visual-antmaze-large-v0 --save_path=data/visual-antmaze-large-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-giant-stitch-v0
python generate_locomaze.py --env_name=visual-antmaze-giant-v0 --save_path=data/visual-antmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-teleport-stitch-v0
python generate_locomaze.py --env_name=visual-antmaze-teleport-v0 --save_path=data/visual-antmaze-teleport-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-medium-explore-v0
python generate_locomaze.py --env_name=visual-antmaze-medium-v0 --save_path=data/visual-antmaze-medium-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-large-explore-v0
python generate_locomaze.py --env_name=visual-antmaze-large-v0 --save_path=data/visual-antmaze-large-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# visual-antmaze-teleport-explore-v0
python generate_locomaze.py --env_name=visual-antmaze-teleport-v0 --save_path=data/visual-antmaze-teleport-explore-v0.npz --dataset_type=explore --num_episodes=10000 --max_episode_steps=501 --noise=1.0 --restore_path=experts/ant --restore_epoch=400000
# visual-humanoidmaze-medium-navigate-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-medium-v0 --save_path=data/visual-humanoidmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=2001 --restore_path=experts/humanoid --restore_epoch=40000000
# visual-humanoidmaze-large-navigate-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-large-v0 --save_path=data/visual-humanoidmaze-large-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=2001 --restore_path=experts/humanoid --restore_epoch=40000000
# visual-humanoidmaze-giant-navigate-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-giant-v0 --save_path=data/visual-humanoidmaze-giant-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=4001 --restore_path=experts/humanoid --restore_epoch=40000000
# visual-humanoidmaze-medium-stitch-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-medium-v0 --save_path=data/visual-humanoidmaze-medium-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# visual-humanoidmaze-large-stitch-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-large-v0 --save_path=data/visual-humanoidmaze-large-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# visual-humanoidmaze-giant-stitch-v0
python generate_locomaze.py --env_name=visual-humanoidmaze-giant-v0 --save_path=data/visual-humanoidmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=10000 --max_episode_steps=401 --restore_path=experts/humanoid --restore_epoch=40000000
# cube-single-play-v0
python generate_manipspace.py --env_name=cube-single-v0 --save_path=data/cube-single-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# cube-double-play-v0
python generate_manipspace.py --env_name=cube-double-v0 --save_path=data/cube-double-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# cube-triple-play-v0
python generate_manipspace.py --env_name=cube-triple-v0 --save_path=data/cube-triple-play-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=play
# cube-quadruple-play-v0
python generate_manipspace.py --env_name=cube-quadruple-v0 --save_path=data/cube-quadruple-play-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=play
# cube-single-noisy-v0
python generate_manipspace.py --env_name=cube-single-v0 --save_path=data/cube-single-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# cube-double-noisy-v0
python generate_manipspace.py --env_name=cube-double-v0 --save_path=data/cube-double-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# cube-triple-noisy-v0
python generate_manipspace.py --env_name=cube-triple-v0 --save_path=data/cube-triple-noisy-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# cube-quadruple-noisy-v0
python generate_manipspace.py --env_name=cube-quadruple-v0 --save_path=data/cube-quadruple-noisy-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# scene-play-v0
python generate_manipspace.py --env_name=scene-v0 --save_path=data/scene-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# scene-noisy-v0
python generate_manipspace.py --env_name=scene-v0 --save_path=data/scene-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# puzzle-3x3-play-v0
python generate_manipspace.py --env_name=puzzle-3x3-v0 --save_path=data/puzzle-3x3-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# puzzle-4x4-play-v0
python generate_manipspace.py --env_name=puzzle-4x4-v0 --save_path=data/puzzle-4x4-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# puzzle-4x5-play-v0
python generate_manipspace.py --env_name=puzzle-4x5-v0 --save_path=data/puzzle-4x5-play-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=play
# puzzle-4x6-play-v0
python generate_manipspace.py --env_name=puzzle-4x6-v0 --save_path=data/puzzle-4x6-play-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=play
# puzzle-3x3-noisy-v0
python generate_manipspace.py --env_name=puzzle-3x3-v0 --save_path=data/puzzle-3x3-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# puzzle-4x4-noisy-v0
python generate_manipspace.py --env_name=puzzle-4x4-v0 --save_path=data/puzzle-4x4-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# puzzle-4x5-noisy-v0
python generate_manipspace.py --env_name=puzzle-4x5-v0 --save_path=data/puzzle-4x5-noisy-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# puzzle-4x6-noisy-v0
python generate_manipspace.py --env_name=puzzle-4x6-v0 --save_path=data/puzzle-4x6-noisy-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# visual-cube-single-play-v0
python generate_manipspace.py --env_name=visual-cube-single-v0 --save_path=data/visual-cube-single-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# visual-cube-double-play-v0
python generate_manipspace.py --env_name=visual-cube-double-v0 --save_path=data/visual-cube-double-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# visual-cube-triple-play-v0
python generate_manipspace.py --env_name=visual-cube-triple-v0 --save_path=data/visual-cube-triple-play-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=play
# visual-cube-quadruple-play-v0
python generate_manipspace.py --env_name=visual-cube-quadruple-v0 --save_path=data/visual-cube-quadruple-play-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=play
# visual-cube-single-noisy-v0
python generate_manipspace.py --env_name=visual-cube-single-v0 --save_path=data/visual-cube-single-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# visual-cube-double-noisy-v0
python generate_manipspace.py --env_name=visual-cube-double-v0 --save_path=data/visual-cube-double-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# visual-cube-triple-noisy-v0
python generate_manipspace.py --env_name=visual-cube-triple-v0 --save_path=data/visual-cube-triple-noisy-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# visual-cube-quadruple-noisy-v0
python generate_manipspace.py --env_name=visual-cube-quadruple-v0 --save_path=data/visual-cube-quadruple-noisy-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# visual-scene-play-v0
python generate_manipspace.py --env_name=visual-scene-v0 --save_path=data/visual-scene-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# visual-scene-noisy-v0
python generate_manipspace.py --env_name=visual-scene-v0 --save_path=data/visual-scene-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.1
# visual-puzzle-3x3-play-v0
python generate_manipspace.py --env_name=visual-puzzle-3x3-v0 --save_path=data/visual-puzzle-3x3-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# visual-puzzle-4x4-play-v0
python generate_manipspace.py --env_name=visual-puzzle-4x4-v0 --save_path=data/visual-puzzle-4x4-play-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=play
# visual-puzzle-4x5-play-v0
python generate_manipspace.py --env_name=visual-puzzle-4x5-v0 --save_path=data/visual-puzzle-4x5-play-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=play
# visual-puzzle-4x6-play-v0
python generate_manipspace.py --env_name=visual-puzzle-4x6-v0 --save_path=data/visual-puzzle-4x6-play-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=play
# visual-puzzle-3x3-noisy-v0
python generate_manipspace.py --env_name=visual-puzzle-3x3-v0 --save_path=data/visual-puzzle-3x3-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# visual-puzzle-4x4-noisy-v0
python generate_manipspace.py --env_name=visual-puzzle-4x4-v0 --save_path=data/visual-puzzle-4x4-noisy-v0.npz --num_episodes=1000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# visual-puzzle-4x5-noisy-v0
python generate_manipspace.py --env_name=visual-puzzle-4x5-v0 --save_path=data/visual-puzzle-4x5-noisy-v0.npz --num_episodes=3000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# visual-puzzle-4x6-noisy-v0
python generate_manipspace.py --env_name=visual-puzzle-4x6-v0 --save_path=data/visual-puzzle-4x6-noisy-v0.npz --num_episodes=5000 --max_episode_steps=1001 --dataset_type=noisy --p_random_action=0.2
# powderworld-easy-play-v0
python generate_powderworld.py --env_name=powderworld-easy-v0 --save_path=data/powderworld-easy-play-v0.npz --dataset_type=play --num_episodes=1000 --max_episode_steps=1001
# powderworld-medium-play-v0
python generate_powderworld.py --env_name=powderworld-medium-v0 --save_path=data/powderworld-medium-play-v0.npz --dataset_type=play --num_episodes=3000 --max_episode_steps=1001
# powderworld-hard-play-v0
python generate_powderworld.py --env_name=powderworld-hard-v0 --save_path=data/powderworld-hard-play-v0.npz --dataset_type=play --num_episodes=5000 --max_episode_steps=1001
================================================
FILE: data_gen_scripts/generate_antsoccer.py
================================================
import glob
import json
import pathlib
from collections import defaultdict
import gymnasium
import numpy as np
from absl import app, flags
from agents import SACAgent
from tqdm import trange
from utils.evaluation import supply_rng
from utils.flax_utils import restore_agent
import ogbench.locomaze # noqa
FLAGS = flags.FLAGS
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'antsoccer-arena-v0', 'Environment name.')
flags.DEFINE_string('dataset_type', 'navigate', 'Dataset type.')
flags.DEFINE_string('loco_restore_path', 'experts/ant', 'Locomotion agent restore path.')
flags.DEFINE_integer('loco_restore_epoch', 400000, 'Locomotion agent restore epoch.')
flags.DEFINE_string('ball_restore_path', 'experts/antball', 'Ball agent restore path.')
flags.DEFINE_integer('ball_restore_epoch', 12000000, 'Ball agent restore epoch.')
flags.DEFINE_string('save_path', None, 'Save path.')
flags.DEFINE_float('noise', 0.2, 'Gaussian action noise level.')
flags.DEFINE_integer('num_episodes', 1000, 'Number of episodes.')
flags.DEFINE_integer('max_episode_steps', 1001, 'Maximum number of steps in an episode.')
def load_agent(restore_path, restore_epoch, ob_dim, action_dim):
"""Initialize and load a SAC agent from a given path."""
# Load agent config.
candidates = glob.glob(restore_path)
assert len(candidates) == 1, f'Found {len(candidates)} candidates: {candidates}'
with open(candidates[0] + '/flags.json', 'r') as f:
agent_config = json.load(f)['agent']
# Load agent.
agent = SACAgent.create(
FLAGS.seed,
np.zeros(ob_dim),
np.zeros(action_dim),
agent_config,
)
agent = restore_agent(agent, restore_path, restore_epoch)
return agent
def main(_):
assert FLAGS.dataset_type in ['navigate', 'stitch']
# 'navigate': Repeatedly navigate to the ball and then to a goal in a single episode.
# 'stitch': Either only navigate or only dribble the ball to a goal in a single episode.
# Initialize environment.
env = gymnasium.make(
FLAGS.env_name,
terminate_at_goal=False,
max_episode_steps=FLAGS.max_episode_steps,
)
ob_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
# Initialize oracle agent.
loco_agent = load_agent(FLAGS.loco_restore_path, FLAGS.loco_restore_epoch, ob_dim, action_dim)
ball_agent = load_agent(FLAGS.ball_restore_path, FLAGS.ball_restore_epoch, ob_dim, action_dim)
loco_actor_fn = supply_rng(loco_agent.sample_actions, rng=loco_agent.rng)
ball_actor_fn = supply_rng(ball_agent.sample_actions, rng=ball_agent.rng)
def get_agent_action(ob, goal_xy):
"""Get an action for the agent to navigate to the goal."""
if 'arena' not in FLAGS.env_name:
# In the actual maze environment, replace the goal with the oracle subgoal.
goal_xy, _ = env.unwrapped.get_oracle_subgoal(ob[:2], goal_xy)
goal_dir = goal_xy - ob[:2]
goal_dir = goal_dir / (np.linalg.norm(goal_dir) + 1e-6)
# Concatenate the agent's joint positions (excluding the x-y position), joint velocities, and goal direction.
agent_ob = np.concatenate([ob[2:15], ob[22:36], goal_dir])
action = loco_actor_fn(agent_ob, temperature=0)
return action
def get_ball_action(ob, ball_xy, goal_xy):
"""Get an action for the agent to dribble the ball to the goal."""
if 'arena' in FLAGS.env_name:
if np.linalg.norm(goal_xy - ball_xy) > 10:
# If the ball is too far from the goal, set a virtual goal 10 units away from the ball. This is because
# the ball agent is not trained to dribble the ball to the goal that is too far away.
goal_xy = ball_xy + 10 * (goal_xy - ball_xy) / np.linalg.norm(goal_xy - ball_xy)
else:
# In the actual maze environment, replace the goal with the oracle subgoal.
goal_xy, _ = env.unwrapped.get_oracle_subgoal(ball_xy, goal_xy)
# Concatenate the agent and ball's joint positions (excluding their x-y positions), their joint velocities, and
# the relative positions of the ball and the goal.
agent_ob = np.concatenate([ob[2:15], ob[17:], ball_xy - agent_xy, goal_xy - ball_xy])
action = ball_actor_fn(agent_ob, temperature=0)
return action
# Store all empty cells.
all_cells = []
maze_map = env.unwrapped.maze_map
for i in range(maze_map.shape[0]):
for j in range(maze_map.shape[1]):
if maze_map[i, j] == 0:
all_cells.append((i, j))
# Collect data.
dataset = defaultdict(list)
total_steps = 0
total_train_steps = 0
num_train_episodes = FLAGS.num_episodes
num_val_episodes = FLAGS.num_episodes // 10
for ep_idx in trange(num_train_episodes + num_val_episodes):
if FLAGS.dataset_type == 'navigate':
# Sample random initial positions for the agent, the ball, and the goal.
agent_init_idx, ball_init_idx, goal_idx = np.random.choice(len(all_cells), 3, replace=False)
agent_init_ij = all_cells[agent_init_idx]
ball_init_ij = all_cells[ball_init_idx]
goal_ij = all_cells[goal_idx]
elif FLAGS.dataset_type == 'stitch':
# Randomly choose between the 'navigate' and 'dribble' modes.
cur_mode = 'navigate' if np.random.randint(2) == 0 else 'dribble'
# Sample random initial positions for the agent, the ball, and the goal. In the 'dribble' mode, the ball
# always starts at the agent's position.
agent_init_idx, ball_init_idx, goal_idx = np.random.choice(len(all_cells), 3, replace=False)
agent_init_ij = all_cells[agent_init_idx]
ball_init_ij = all_cells[ball_init_idx] if cur_mode == 'navigate' else agent_init_ij
goal_ij = all_cells[goal_idx]
else:
raise ValueError(f'Unsupported dataset_type: {FLAGS.dataset_type}')
ob, _ = env.reset(
options=dict(task_info=dict(agent_init_ij=agent_init_ij, ball_init_ij=ball_init_ij, goal_ij=goal_ij))
)
done = False
step = 0
virtual_agent_goal_xy = None # Virtual goal for the agent to move to when stuck.
while not done:
agent_xy, ball_xy = env.unwrapped.get_agent_ball_xy()
agent_xy, ball_xy = np.array(agent_xy), np.array(ball_xy)
goal_xy = np.array(env.unwrapped.cur_goal_xy)
if FLAGS.dataset_type == 'navigate':
if virtual_agent_goal_xy is None:
if np.linalg.norm(agent_xy - ball_xy) > 2:
# If the agent is far from the ball, move to the ball.
action = get_agent_action(ob, ball_xy)
else:
# If the agent is close to the ball, dribble the ball to the goal.
action = get_ball_action(ob, ball_xy, goal_xy)
else:
# When virtual_agent_goal_xy is set, move to the virtual goal.
action = get_agent_action(ob, virtual_agent_goal_xy)
elif FLAGS.dataset_type == 'stitch':
if cur_mode == 'navigate':
# Navigate to the goal.
action = get_agent_action(ob, goal_xy)
else:
# Dribble the ball to the goal.
action = get_ball_action(ob, ball_xy, goal_xy)
# Add Gaussian noise to the action.
action = action + np.random.normal(0, FLAGS.noise, action.shape)
action = np.clip(action, -1, 1)
next_ob, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
success = info['success']
if virtual_agent_goal_xy is not None and np.linalg.norm(virtual_agent_goal_xy - next_ob[:2]) <= 0.5:
# If the agent reaches the virtual goal, clear it.
virtual_agent_goal_xy = None
if FLAGS.dataset_type == 'navigate':
if success:
# Sample a new goal state when the current goal is reached.
goal_ij = all_cells[np.random.randint(len(all_cells))]
env.unwrapped.set_goal(goal_ij)
# Determine whether the agent is stuck.
if (
step > 150
and virtual_agent_goal_xy is None
and np.linalg.norm(np.array(dataset['observations'][-150:])[:, :2] - next_ob[:2], axis=1).max() <= 2
):
# When the agent is stuck for 150 steps, set a virtual goal to move to a random cell.
virtual_agent_goal_ij = all_cells[np.random.randint(len(all_cells))]
virtual_agent_goal_xy = np.array(env.unwrapped.ij_to_xy(virtual_agent_goal_ij))
dataset['observations'].append(ob)
dataset['actions'].append(action)
dataset['terminals'].append(done)
dataset['qpos'].append(info['prev_qpos'])
dataset['qvel'].append(info['prev_qvel'])
ob = next_ob
step += 1
total_steps += step
if ep_idx < num_train_episodes:
total_train_steps += step
print('Total steps:', total_steps)
train_path = FLAGS.save_path
val_path = FLAGS.save_path.replace('.npz', '-val.npz')
pathlib.Path(train_path).parent.mkdir(parents=True, exist_ok=True)
# Split the dataset into training and validation sets.
train_dataset = {
k: np.array(v[:total_train_steps], dtype=np.float32 if k != 'terminals' else bool) for k, v in dataset.items()
}
val_dataset = {
k: np.array(v[total_train_steps:], dtype=np.float32 if k != 'terminals' else bool) for k, v in dataset.items()
}
for path, dataset in [(train_path, train_dataset), (val_path, val_dataset)]:
np.savez_compressed(path, **dataset)
if __name__ == '__main__':
app.run(main)
================================================
FILE: data_gen_scripts/generate_locomaze.py
================================================
import glob
import json
import pathlib
from collections import defaultdict
import gymnasium
import numpy as np
from absl import app, flags
from agents import SACAgent
from tqdm import trange
from utils.evaluation import supply_rng
from utils.flax_utils import restore_agent
import ogbench.locomaze # noqa
FLAGS = flags.FLAGS
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'antmaze-large-v0', 'Environment name.')
flags.DEFINE_string('dataset_type', 'navigate', 'Dataset type.')
flags.DEFINE_string('restore_path', 'experts/ant', 'Expert agent restore path.')
flags.DEFINE_integer('restore_epoch', 400000, 'Expert agent restore epoch.')
flags.DEFINE_string('save_path', None, 'Save path.')
flags.DEFINE_float('noise', 0.2, 'Gaussian action noise level.')
flags.DEFINE_integer('num_episodes', 1000, 'Number of episodes.')
flags.DEFINE_integer('max_episode_steps', 1001, 'Maximum number of steps in an episode.')
def main(_):
assert FLAGS.dataset_type in ['path', 'navigate', 'stitch', 'explore']
# 'path': Reach a single goal and stay there.
# 'navigate': Repeatedly reach randomly sampled goals in a single episode.
# 'stitch': Reach a nearby goal that is 4 cells away and stay there.
# 'explore': Repeatedly follow random directions sampled every 10 steps.
# Initialize environment.
env = gymnasium.make(
FLAGS.env_name,
terminate_at_goal=False,
max_episode_steps=FLAGS.max_episode_steps,
)
ob_dim = env.observation_space.shape[0]
# Initialize oracle agent.
if 'point' in FLAGS.env_name:
def actor_fn(ob, temperature):
return ob[-2:]
else:
# Load agent config.
restore_path = FLAGS.restore_path
candidates = glob.glob(restore_path)
assert len(candidates) == 1, f'Found {len(candidates)} candidates: {candidates}'
with open(candidates[0] + '/flags.json', 'r') as f:
agent_config = json.load(f)['agent']
# Load agent.
agent = SACAgent.create(
FLAGS.seed,
np.zeros(ob_dim),
env.action_space.sample(),
agent_config,
)
agent = restore_agent(agent, FLAGS.restore_path, FLAGS.restore_epoch)
actor_fn = supply_rng(agent.sample_actions, rng=agent.rng)
# Store all empty cells and vertex cells.
all_cells = []
vertex_cells = []
maze_map = env.unwrapped.maze_map
for i in range(maze_map.shape[0]):
for j in range(maze_map.shape[1]):
if maze_map[i, j] == 0:
all_cells.append((i, j))
# Exclude hallway cells.
if (
maze_map[i - 1, j] == 0
and maze_map[i + 1, j] == 0
and maze_map[i, j - 1] == 1
and maze_map[i, j + 1] == 1
):
continue
if (
maze_map[i, j - 1] == 0
and maze_map[i, j + 1] == 0
and maze_map[i - 1, j] == 1
and maze_map[i + 1, j] == 1
):
continue
vertex_cells.append((i, j))
# Collect data.
dataset = defaultdict(list)
total_steps = 0
total_train_steps = 0
num_train_episodes = FLAGS.num_episodes
num_val_episodes = FLAGS.num_episodes // 10
for ep_idx in trange(num_train_episodes + num_val_episodes):
if FLAGS.dataset_type in ['path', 'navigate', 'explore']:
# Sample an initial state from all cells.
init_ij = all_cells[np.random.randint(len(all_cells))]
# Sample a goal state from vertex cells.
goal_ij = vertex_cells[np.random.randint(len(vertex_cells))]
elif FLAGS.dataset_type == 'stitch':
# Sample an initial state from all cells.
init_ij = all_cells[np.random.randint(len(all_cells))]
# Perform BFS to find adjacent cells.
adj_cells = []
adj_steps = 4 # Target distance from the initial cell.
bfs_map = maze_map.copy()
for i in range(bfs_map.shape[0]):
for j in range(bfs_map.shape[1]):
bfs_map[i][j] = -1
bfs_map[init_ij[0], init_ij[1]] = 0
queue = [init_ij]
while len(queue) > 0:
i, j = queue.pop(0)
for di, dj in [(-1, 0), (0, -1), (1, 0), (0, 1)]:
ni, nj = i + di, j + dj
if (
0 <= ni < bfs_map.shape[0]
and 0 <= nj < bfs_map.shape[1]
and maze_map[ni, nj] == 0
and bfs_map[ni, nj] == -1
):
bfs_map[ni][nj] = bfs_map[i][j] + 1
queue.append((ni, nj))
if bfs_map[ni][nj] == adj_steps:
adj_cells.append((ni, nj))
# Sample a goal state from adjacent cells.
goal_ij = adj_cells[np.random.randint(len(adj_cells))] if len(adj_cells) > 0 else init_ij
else:
raise ValueError(f'Unsupported dataset_type: {FLAGS.dataset_type}')
ob, _ = env.reset(options=dict(task_info=dict(init_ij=init_ij, goal_ij=goal_ij)))
done = False
step = 0
cur_subgoal_dir = None # Current subgoal direction (only for 'explore').
while not done:
if FLAGS.dataset_type == 'explore':
# Sample a random direction every 10 steps.
if step % 10 == 0:
cur_subgoal_dir = np.random.randn(2)
cur_subgoal_dir = cur_subgoal_dir / (np.linalg.norm(cur_subgoal_dir) + 1e-6)
subgoal_dir = cur_subgoal_dir
else:
# Get the oracle subgoal and compute the direction.
subgoal_xy, _ = env.unwrapped.get_oracle_subgoal(env.unwrapped.get_xy(), env.unwrapped.cur_goal_xy)
subgoal_dir = subgoal_xy - env.unwrapped.get_xy()
subgoal_dir = subgoal_dir / (np.linalg.norm(subgoal_dir) + 1e-6)
agent_ob = env.unwrapped.get_ob(ob_type='states')
# Exclude the agent's position and add the subgoal direction.
agent_ob = np.concatenate([agent_ob[2:], subgoal_dir])
action = actor_fn(agent_ob, temperature=0)
# Add Gaussian noise to the action.
action = action + np.random.normal(0, FLAGS.noise, action.shape)
action = np.clip(action, -1, 1)
next_ob, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
success = info['success']
# Sample a new goal state when the current goal is reached.
if success and FLAGS.dataset_type == 'navigate':
goal_ij = vertex_cells[np.random.randint(len(vertex_cells))]
env.unwrapped.set_goal(goal_ij)
dataset['observations'].append(ob)
dataset['actions'].append(action)
dataset['terminals'].append(done)
dataset['qpos'].append(info['prev_qpos'])
dataset['qvel'].append(info['prev_qvel'])
ob = next_ob
step += 1
total_steps += step
if ep_idx < num_train_episodes:
total_train_steps += step
print('Total steps:', total_steps)
train_path = FLAGS.save_path
val_path = FLAGS.save_path.replace('.npz', '-val.npz')
pathlib.Path(train_path).parent.mkdir(parents=True, exist_ok=True)
# Split the dataset into training and validation sets.
train_dataset = {}
val_dataset = {}
for k, v in dataset.items():
if 'observations' in k and v[0].dtype == np.uint8:
dtype = np.uint8
elif k == 'terminals':
dtype = bool
else:
dtype = np.float32
train_dataset[k] = np.array(v[:total_train_steps], dtype=dtype)
val_dataset[k] = np.array(v[total_train_steps:], dtype=dtype)
for path, dataset in [(train_path, train_dataset), (val_path, val_dataset)]:
np.savez_compressed(path, **dataset)
if __name__ == '__main__':
app.run(main)
================================================
FILE: data_gen_scripts/generate_manipspace.py
================================================
import pathlib
from collections import defaultdict
import gymnasium
import numpy as np
from absl import app, flags
from tqdm import trange
import ogbench.manipspace # noqa
from ogbench.manipspace.oracles.markov.button_markov import ButtonMarkovOracle
from ogbench.manipspace.oracles.markov.cube_markov import CubeMarkovOracle
from ogbench.manipspace.oracles.markov.drawer_markov import DrawerMarkovOracle
from ogbench.manipspace.oracles.markov.window_markov import WindowMarkovOracle
from ogbench.manipspace.oracles.plan.button_plan import ButtonPlanOracle
from ogbench.manipspace.oracles.plan.cube_plan import CubePlanOracle
from ogbench.manipspace.oracles.plan.drawer_plan import DrawerPlanOracle
from ogbench.manipspace.oracles.plan.window_plan import WindowPlanOracle
FLAGS = flags.FLAGS
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'cube-single-v0', 'Environment name.')
flags.DEFINE_string('dataset_type', 'play', 'Dataset type.')
flags.DEFINE_string('save_path', None, 'Save path.')
flags.DEFINE_float('noise', 0.1, 'Action noise level.')
flags.DEFINE_float('noise_smoothing', 0.5, 'Action noise smoothing level for PlanOracle.')
flags.DEFINE_float('min_norm', 0.4, 'Minimum action norm for MarkovOracle.')
flags.DEFINE_float('p_random_action', 0, 'Probability of selecting a random action.')
flags.DEFINE_integer('num_episodes', 1000, 'Number of episodes.')
flags.DEFINE_integer('max_episode_steps', 1001, 'Number of episodes.')
def main(_):
assert FLAGS.dataset_type in ['play', 'noisy']
# 'play': Use a non-Markovian oracle (PlanOracle) that follows a pre-computed plan.
# 'noisy': Use a Markovian, closed-loop oracle (MarkovOracle) with Gaussian action noise.
# Initialize environment.
env = gymnasium.make(
FLAGS.env_name,
terminate_at_goal=False,
mode='data_collection',
max_episode_steps=FLAGS.max_episode_steps,
)
# Initialize oracles.
oracle_type = 'plan' if FLAGS.dataset_type == 'play' else 'markov'
has_button_states = hasattr(env.unwrapped, '_cur_button_states')
if 'cube' in FLAGS.env_name:
if oracle_type == 'markov':
agents = {
'cube': CubeMarkovOracle(env=env, min_norm=FLAGS.min_norm),
}
else:
agents = {
'cube': CubePlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
}
elif 'scene' in FLAGS.env_name:
if oracle_type == 'markov':
agents = {
'cube': CubeMarkovOracle(env=env, min_norm=FLAGS.min_norm, max_step=100),
'button': ButtonMarkovOracle(env=env, min_norm=FLAGS.min_norm),
'drawer': DrawerMarkovOracle(env=env, min_norm=FLAGS.min_norm),
'window': WindowMarkovOracle(env=env, min_norm=FLAGS.min_norm),
}
else:
agents = {
'cube': CubePlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
'button': ButtonPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
'drawer': DrawerPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
'window': WindowPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
}
elif 'puzzle' in FLAGS.env_name:
if oracle_type == 'markov':
agents = {
'button': ButtonMarkovOracle(env=env, min_norm=FLAGS.min_norm, gripper_always_closed=True),
}
else:
agents = {
'button': ButtonPlanOracle(
env=env,
noise=FLAGS.noise,
noise_smoothing=FLAGS.noise_smoothing,
gripper_always_closed=True,
),
}
# Collect data.
dataset = defaultdict(list)
total_steps = 0
total_train_steps = 0
num_train_episodes = FLAGS.num_episodes
num_val_episodes = FLAGS.num_episodes // 10
for ep_idx in trange(num_train_episodes + num_val_episodes):
# Have an additional while loop to handle rare cases with undesirable states (for the Scene environment).
while True:
ob, info = env.reset()
# Set the cube stacking probability for this episode.
if 'single' in FLAGS.env_name:
p_stack = 0.0
elif 'double' in FLAGS.env_name:
p_stack = np.random.uniform(0.0, 0.25)
elif 'triple' in FLAGS.env_name:
p_stack = np.random.uniform(0.05, 0.35)
elif 'quadruple' in FLAGS.env_name:
p_stack = np.random.uniform(0.1, 0.5)
elif 'octuple' in FLAGS.env_name:
p_stack = np.random.uniform(0.0, 0.35)
else:
p_stack = 0.5
if oracle_type == 'markov':
# Set the action noise level for this episode.
xi = np.random.uniform(0, FLAGS.noise)
agent = agents[info['privileged/target_task']]
agent.reset(ob, info)
done = False
step = 0
ep_qpos = []
while not done:
if np.random.rand() < FLAGS.p_random_action:
# Sample a random action.
action = env.action_space.sample()
else:
# Get an action from the oracle.
action = agent.select_action(ob, info)
action = np.array(action)
if oracle_type == 'markov':
# Add Gaussian noise to the action.
action = action + np.random.normal(0, [xi, xi, xi, xi * 3, xi * 10], action.shape)
action = np.clip(action, -1, 1)
next_ob, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
if agent.done:
# Set a new task when the current task is done.
agent_ob, agent_info = env.unwrapped.set_new_target(p_stack=p_stack)
agent = agents[agent_info['privileged/target_task']]
agent.reset(agent_ob, agent_info)
dataset['observations'].append(ob)
dataset['actions'].append(action)
dataset['terminals'].append(done)
dataset['qpos'].append(info['prev_qpos'])
dataset['qvel'].append(info['prev_qvel'])
if has_button_states:
dataset['button_states'].append(info['prev_button_states'])
ep_qpos.append(info['prev_qpos'])
ob = next_ob
step += 1
if 'scene' in FLAGS.env_name:
# Perform health check. We want to ensure that the cube is always visible unless it's in the drawer.
# Otherwise, the test-time goal images may become ambiguous.
is_healthy = True
ep_qpos = np.array(ep_qpos)
block_xyzs = ep_qpos[:, 14:17]
if (block_xyzs[:, 1] >= 0.29).any():
is_healthy = False # Block goes too far right.
if ((block_xyzs[:, 1] <= -0.3) & ((block_xyzs[:, 2] < 0.06) | (block_xyzs[:, 2] > 0.08))).any():
is_healthy = False # Block goes too far left, without being in the drawer.
if is_healthy:
break
else:
# Remove the last episode and retry.
print('Unhealthy episode, retrying...', flush=True)
for k in dataset.keys():
dataset[k] = dataset[k][:-step]
else:
break
total_steps += step
if ep_idx < num_train_episodes:
total_train_steps += step
print('Total steps:', total_steps)
train_path = FLAGS.save_path
val_path = FLAGS.save_path.replace('.npz', '-val.npz')
pathlib.Path(train_path).parent.mkdir(parents=True, exist_ok=True)
# Split the dataset into training and validation sets.
train_dataset = {}
val_dataset = {}
for k, v in dataset.items():
if 'observations' in k and v[0].dtype == np.uint8:
dtype = np.uint8
elif k == 'terminals':
dtype = bool
elif k == 'button_states':
dtype = np.int64
else:
dtype = np.float32
train_dataset[k] = np.array(v[:total_train_steps], dtype=dtype)
val_dataset[k] = np.array(v[total_train_steps:], dtype=dtype)
for path, dataset in [(train_path, train_dataset), (val_path, val_dataset)]:
np.savez_compressed(path, **dataset)
if __name__ == '__main__':
app.run(main)
================================================
FILE: data_gen_scripts/generate_powderworld.py
================================================
import pathlib
from collections import defaultdict
import gymnasium
import numpy as np
from absl import app, flags
from tqdm import trange
import ogbench.powderworld # noqa
from ogbench.powderworld.behaviors import FillBehavior, LineBehavior, SquareBehavior
FLAGS = flags.FLAGS
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'powderworld-v0', 'Environment name.')
flags.DEFINE_string('dataset_type', 'play', 'Dataset type.')
flags.DEFINE_string('save_path', None, 'Save path.')
flags.DEFINE_integer('num_episodes', 1000, 'Number of episodes.')
flags.DEFINE_integer('max_episode_steps', 1001, 'Maximum number of steps in an episode.')
flags.DEFINE_float('p_random_action', 0.5, 'Probability of selecting a random action.')
def main(_):
assert FLAGS.dataset_type in ['play']
# Initialize environment.
env = gymnasium.make(
FLAGS.env_name,
mode='data_collection',
max_episode_steps=FLAGS.max_episode_steps,
)
env.reset()
# Initialize agents.
agents = [
FillBehavior(env=env),
LineBehavior(env=env),
SquareBehavior(env=env),
]
probs = np.array([1, 3, 3]) # Agent selection probabilities.
probs = probs / probs.sum()
# Collect data.
dataset = defaultdict(list)
total_steps = 0
total_train_steps = 0
num_train_episodes = FLAGS.num_episodes
num_val_episodes = FLAGS.num_episodes // 10
for ep_idx in trange(num_train_episodes + num_val_episodes):
ob, info = env.reset()
agent = np.random.choice(agents, p=probs)
agent.reset(ob, info)
done = False
step = 0
action_step = 0 # Action cycle counter (0, 1, 2).
while not done:
if action_step == 0:
# Select an action every 3 steps.
if np.random.rand() < FLAGS.p_random_action:
# Sample a random action.
semantic_action = env.unwrapped.sample_semantic_action()
else:
# Get an action from the agent.
semantic_action = agent.select_action(ob, info)
action = env.unwrapped.semantic_action_to_action(*semantic_action)
next_ob, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
if agent.done and FLAGS.dataset_type == 'play':
agent = np.random.choice(agents, p=probs)
agent.reset(ob, info)
dataset['observations'].append(ob)
dataset['actions'].append(action)
dataset['terminals'].append(done)
ob = next_ob
step += 1
action_step = (action_step + 1) % 3
total_steps += step
if ep_idx < num_train_episodes:
total_train_steps += step
print('Total steps:', total_steps)
train_path = FLAGS.save_path
val_path = FLAGS.save_path.replace('.npz', '-val.npz')
pathlib.Path(train_path).parent.mkdir(parents=True, exist_ok=True)
# Split the dataset into training and validation sets.
train_dataset = {}
val_dataset = {}
for k, v in dataset.items():
if 'observations' in k and v[0].dtype == np.uint8:
dtype = np.uint8
elif 'actions':
dtype = np.int32
elif k == 'terminals':
dtype = bool
else:
dtype = np.float32
train_dataset[k] = np.array(v[:total_train_steps], dtype=dtype)
val_dataset[k] = np.array(v[total_train_steps:], dtype=dtype)
for path, dataset in [(train_path, train_dataset), (val_path, val_dataset)]:
np.savez_compressed(path, **dataset)
if __name__ == '__main__':
app.run(main)
================================================
FILE: data_gen_scripts/main_sac.py
================================================
import json
import os
import random
import time
import jax
import numpy as np
import tqdm
import wandb
from absl import app, flags
from agents import agents
from ml_collections import config_flags
from online_env_utils import make_online_env
from utils.datasets import ReplayBuffer
from utils.evaluation import evaluate, flatten
from utils.flax_utils import restore_agent, save_agent
from utils.log_utils import CsvLogger, get_exp_name, get_flag_dict, get_wandb_video, setup_wandb
from viz_utils import visualize_trajs
FLAGS = flags.FLAGS
flags.DEFINE_string('run_group', 'Debug', 'Run group.')
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'online-ant-xy-v0', 'Environment name.')
flags.DEFINE_string('save_dir', 'exp/', 'Save directory.')
flags.DEFINE_string('restore_path', None, 'Restore path.')
flags.DEFINE_integer('restore_epoch', None, 'Restore epoch.')
flags.DEFINE_integer('seed_steps', 10000, 'Number of seed steps.')
flags.DEFINE_integer('train_steps', 1000000, 'Number of training steps.')
flags.DEFINE_integer('train_interval', 1, 'Train interval.')
flags.DEFINE_integer('num_epochs', 1, 'Number of updates per train interval.')
flags.DEFINE_integer('log_interval', 5000, 'Logging interval.')
flags.DEFINE_integer('eval_interval', 100000, 'Evaluation interval.')
flags.DEFINE_integer('save_interval', 1000000, 'Saving interval.')
flags.DEFINE_integer('reset_interval', 0, 'Full parameter reset interval.')
flags.DEFINE_integer('terminate_at_end', 0, 'Whether to set terminated=True when truncated=True.')
flags.DEFINE_integer('eval_episodes', 50, 'Number of episodes for each task.')
flags.DEFINE_float('eval_temperature', 0, 'Actor temperature for evaluation.')
flags.DEFINE_float('eval_gaussian', None, 'Action Gaussian noise for evaluation.')
flags.DEFINE_integer('video_episodes', 1, 'Number of video episodes for each task.')
flags.DEFINE_integer('video_frame_skip', 3, 'Frame skip for videos.')
flags.DEFINE_integer('eval_on_cpu', 1, 'Whether to evaluate on CPU.')
config_flags.DEFINE_config_file('agent', '../impls/agents/sac.py', lock_config=False)
def main(_):
# Set up logger.
exp_name = get_exp_name(FLAGS.seed)
setup_wandb(project='OGBench', group=FLAGS.run_group, name=exp_name)
FLAGS.save_dir = os.path.join(FLAGS.save_dir, wandb.run.project, FLAGS.run_group, exp_name)
os.makedirs(FLAGS.save_dir, exist_ok=True)
flag_dict = get_flag_dict()
with open(os.path.join(FLAGS.save_dir, 'flags.json'), 'w') as f:
json.dump(flag_dict, f)
config = FLAGS.agent
# Set up environments and replay buffer.
env = make_online_env(FLAGS.env_name)
eval_env = make_online_env(FLAGS.env_name)
example_transition = dict(
observations=env.observation_space.sample(),
actions=env.action_space.sample(),
rewards=0.0,
masks=1.0,
next_observations=env.observation_space.sample(),
)
replay_buffer = ReplayBuffer.create(example_transition, size=int(1e6))
# Initialize agent.
random.seed(FLAGS.seed)
np.random.seed(FLAGS.seed)
agent_class = agents[config['agent_name']]
agent = agent_class.create(
FLAGS.seed,
example_transition['observations'],
example_transition['actions'],
config,
)
# Restore agent.
if FLAGS.restore_path is not None:
agent = restore_agent(agent, FLAGS.restore_path, FLAGS.restore_epoch)
# Train agent.
expl_metrics = dict()
expl_rng = jax.random.PRNGKey(FLAGS.seed)
ob, _ = env.reset()
train_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'train.csv'))
eval_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'eval.csv'))
first_time = time.time()
last_time = time.time()
update_info = None
for i in tqdm.tqdm(range(1, FLAGS.train_steps + 1), smoothing=0.1, dynamic_ncols=True):
# Sample transition.
if i < FLAGS.seed_steps:
action = env.action_space.sample()
else:
expl_rng, key = jax.random.split(expl_rng)
action = agent.sample_actions(observations=ob, seed=key)
action = np.array(action)
next_ob, reward, terminated, truncated, info = env.step(action)
if FLAGS.terminate_at_end and truncated:
terminated = True
replay_buffer.add_transition(
dict(
observations=ob,
actions=action,
rewards=reward,
masks=float(not terminated),
next_observations=next_ob,
)
)
ob = next_ob
if terminated or truncated:
expl_metrics = {f'exploration/{k}': np.mean(v) for k, v in flatten(info).items()}
ob, _ = env.reset()
if replay_buffer.size < FLAGS.seed_steps:
continue
# Update agent.
if i % FLAGS.train_interval == 0:
for _ in range(FLAGS.num_epochs):
batch = replay_buffer.sample(config['batch_size'])
agent, update_info = agent.update(batch)
# Log metrics.
if i % FLAGS.log_interval == 0 and update_info is not None:
train_metrics = {f'training/{k}': v for k, v in update_info.items()}
train_metrics['time/epoch_time'] = (time.time() - last_time) / FLAGS.log_interval
train_metrics['time/total_time'] = time.time() - first_time
train_metrics.update(expl_metrics)
last_time = time.time()
wandb.log(train_metrics, step=i)
train_logger.log(train_metrics, step=i)
# Evaluate agent.
if i % FLAGS.eval_interval == 0:
if FLAGS.eval_on_cpu:
eval_agent = jax.device_put(agent, device=jax.devices('cpu')[0])
else:
eval_agent = agent
eval_metrics = {}
eval_info, trajs, renders = evaluate(
agent=eval_agent,
env=eval_env,
task_id=None,
config=config,
num_eval_episodes=FLAGS.eval_episodes,
num_video_episodes=FLAGS.video_episodes,
video_frame_skip=FLAGS.video_frame_skip,
eval_temperature=FLAGS.eval_temperature,
eval_gaussian=FLAGS.eval_gaussian,
)
eval_metrics.update({f'evaluation/{k}': v for k, v in eval_info.items()})
if FLAGS.video_episodes > 0:
video = get_wandb_video(renders=renders)
eval_metrics['video'] = video
traj_image = visualize_trajs(FLAGS.env_name, trajs)
if traj_image is not None:
eval_metrics['traj'] = wandb.Image(traj_image)
wandb.log(eval_metrics, step=i)
eval_logger.log(eval_metrics, step=i)
# Save agent.
if i % FLAGS.save_interval == 0:
save_agent(agent, FLAGS.save_dir, i)
# Reset agent.
if FLAGS.reset_interval > 0 and i % FLAGS.reset_interval == 0:
new_agent = agent_class.create(
FLAGS.seed + i,
example_transition['observations'],
example_transition['actions'],
config,
)
agent = agent.replace(
network=agent.network.replace(params=new_agent.network.params, opt_state=new_agent.network.opt_state)
)
del new_agent
train_logger.close()
eval_logger.close()
if __name__ == '__main__':
app.run(main)
================================================
FILE: data_gen_scripts/online_env_utils.py
================================================
import gymnasium
from utils.env_utils import EpisodeMonitor
def make_online_env(env_name):
"""Make online environment.
If the environment name contains the '-xy' suffix, the environment will be wrapped with a directional locomotion
wrapper. For example, 'online-ant-xy-v0' will return an 'online-ant-v0' environment wrapped with GymXYWrapper.
Args:
env_name: Name of the environment.
"""
import ogbench.online_locomotion # noqa
# Manually recognize the '-xy' suffix, which indicates that the environment should be wrapped with a directional
# locomotion wrapper.
if '-xy' in env_name:
env_name = env_name.replace('-xy', '')
apply_xy_wrapper = True
else:
apply_xy_wrapper = False
# Set camera.
if 'humanoid' in env_name:
extra_kwargs = dict(camera_id=0)
else:
extra_kwargs = dict()
# Make environment.
env = gymnasium.make(env_name, render_mode='rgb_array', height=200, width=200, **extra_kwargs)
if apply_xy_wrapper:
# Apply the directional locomotion wrapper.
from ogbench.online_locomotion.wrappers import DMCHumanoidXYWrapper, GymXYWrapper
if 'humanoid' in env_name:
env = DMCHumanoidXYWrapper(env, resample_interval=200)
else:
env = GymXYWrapper(env, resample_interval=100)
env = EpisodeMonitor(env)
return env
================================================
FILE: data_gen_scripts/viz_utils.py
================================================
import matplotlib
import numpy as np
from matplotlib import figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
def get_2d_colors(points, min_point, max_point):
"""Get colors corresponding to 2-D points."""
points = np.array(points)
min_point = np.array(min_point)
max_point = np.array(max_point)
colors = (points - min_point) / (max_point - min_point)
colors = np.hstack((colors, (2 - np.sum(colors, axis=1, keepdims=True)) / 2))
colors = np.clip(colors, 0, 1)
colors = np.c_[colors, np.full(len(colors), 0.8)]
return colors
def visualize_trajs(env_name, trajs):
"""Visualize x-y trajectories in locomotion environments.
It reads 'xy' and 'direction' from the 'info' field of the trajectories.
"""
matplotlib.use('Agg')
fig = figure.Figure(tight_layout=True)
canvas = FigureCanvas(fig)
if 'xy' in trajs[0]['info'][0]:
ax = fig.add_subplot()
max_xy = 0.0
for traj in trajs:
xy = np.array([info['xy'] for info in traj['info']])
direction = np.array([info['direction'] for info in traj['info']])
color = get_2d_colors(direction, [-1, -1], [1, 1])
for i in range(len(xy) - 1):
ax.plot(xy[i : i + 2, 0], xy[i : i + 2, 1], color=color[i], linewidth=0.7)
max_xy = max(max_xy, np.abs(xy).max() * 1.2)
plot_axis = [-max_xy, max_xy, -max_xy, max_xy]
ax.axis(plot_axis)
ax.set_aspect('equal')
else:
return None
fig.tight_layout()
canvas.draw()
out_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
out_image = out_image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
return out_image
================================================
FILE: impls/agents/__init__.py
================================================
from agents.crl import CRLAgent
from agents.gcbc import GCBCAgent
from agents.gciql import GCIQLAgent
from agents.gcivl import GCIVLAgent
from agents.hiql import HIQLAgent
from agents.qrl import QRLAgent
from agents.sac import SACAgent
agents = dict(
crl=CRLAgent,
gcbc=GCBCAgent,
gciql=GCIQLAgent,
gcivl=GCIVLAgent,
hiql=HIQLAgent,
qrl=QRLAgent,
sac=SACAgent,
)
================================================
FILE: impls/agents/crl.py
================================================
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import GCActor, GCBilinearValue, GCDiscreteActor, GCDiscreteBilinearCritic
class CRLAgent(flax.struct.PyTreeNode):
"""Contrastive RL (CRL) agent.
This implementation supports both AWR (actor_loss='awr') and DDPG+BC (actor_loss='ddpgbc') for the actor loss.
CRL with DDPG+BC only fits a Q function, while CRL with AWR fits both Q and V functions to compute advantages.
"""
rng: Any
network: Any
config: Any = nonpytree_field()
def contrastive_loss(self, batch, grad_params, module_name='critic'):
"""Compute the contrastive value loss for the Q or V function."""
batch_size = batch['observations'].shape[0]
if module_name == 'critic':
actions = batch['actions']
else:
actions = None
v, phi, psi = self.network.select(module_name)(
batch['observations'],
batch['value_goals'],
actions=actions,
info=True,
params=grad_params,
)
if len(phi.shape) == 2: # Non-ensemble.
phi = phi[None, ...]
psi = psi[None, ...]
logits = jnp.einsum('eik,ejk->ije', phi, psi) / jnp.sqrt(phi.shape[-1])
# logits.shape is (B, B, e) with one term for positive pair and (B - 1) terms for negative pairs in each row.
I = jnp.eye(batch_size)
contrastive_loss = jax.vmap(
lambda _logits: optax.sigmoid_binary_cross_entropy(logits=_logits, labels=I),
in_axes=-1,
out_axes=-1,
)(logits)
contrastive_loss = jnp.mean(contrastive_loss)
# Compute additional statistics.
v = jnp.exp(v)
logits = jnp.mean(logits, axis=-1)
correct = jnp.argmax(logits, axis=1) == jnp.argmax(I, axis=1)
logits_pos = jnp.sum(logits * I) / jnp.sum(I)
logits_neg = jnp.sum(logits * (1 - I)) / jnp.sum(1 - I)
return contrastive_loss, {
'contrastive_loss': contrastive_loss,
'v_mean': v.mean(),
'v_max': v.max(),
'v_min': v.min(),
'binary_accuracy': jnp.mean((logits > 0) == I),
'categorical_accuracy': jnp.mean(correct),
'logits_pos': logits_pos,
'logits_neg': logits_neg,
'logits': logits.mean(),
}
def actor_loss(self, batch, grad_params, rng=None):
"""Compute the actor loss (AWR or DDPG+BC)."""
if self.config['actor_loss'] == 'awr':
# AWR loss.
v = self.network.select('value')(batch['observations'], batch['actor_goals'])
q1, q2 = self.network.select('critic')(batch['observations'], batch['actor_goals'], batch['actions'])
q = jnp.minimum(q1, q2)
adv = q - v
exp_a = jnp.exp(adv * self.config['alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -(exp_a * log_prob).mean()
actor_info = {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
elif self.config['actor_loss'] == 'ddpgbc':
# DDPG+BC loss.
assert not self.config['discrete']
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
if self.config['const_std']:
q_actions = jnp.clip(dist.mode(), -1, 1)
else:
q_actions = jnp.clip(dist.sample(seed=rng), -1, 1)
q1, q2 = self.network.select('critic')(batch['observations'], batch['actor_goals'], q_actions)
q = jnp.minimum(q1, q2)
# Normalize Q values by the absolute mean to make the loss scale invariant.
q_loss = -q.mean() / jax.lax.stop_gradient(jnp.abs(q).mean() + 1e-6)
log_prob = dist.log_prob(batch['actions'])
bc_loss = -(self.config['alpha'] * log_prob).mean()
actor_loss = q_loss + bc_loss
return actor_loss, {
'actor_loss': actor_loss,
'q_loss': q_loss,
'bc_loss': bc_loss,
'q_mean': q.mean(),
'q_abs_mean': jnp.abs(q).mean(),
'bc_log_prob': log_prob.mean(),
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
else:
raise ValueError(f'Unsupported actor loss: {self.config["actor_loss"]}')
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
critic_loss, critic_info = self.contrastive_loss(batch, grad_params, 'critic')
for k, v in critic_info.items():
info[f'critic/{k}'] = v
if self.config['actor_loss'] == 'awr':
value_loss, value_info = self.contrastive_loss(batch, grad_params, 'value')
for k, v in value_info.items():
info[f'value/{k}'] = v
else:
value_loss = 0.0
rng, actor_rng = jax.random.split(rng)
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = critic_loss + value_loss + actor_loss
return loss, info
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define encoders.
encoders = dict()
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
encoders['critic_state'] = encoder_module()
encoders['critic_goal'] = encoder_module()
encoders['actor'] = GCEncoder(concat_encoder=encoder_module())
if config['actor_loss'] == 'awr':
encoders['value_state'] = encoder_module()
encoders['value_goal'] = encoder_module()
# Define value and actor networks.
if config['discrete']:
critic_def = GCDiscreteBilinearCritic(
hidden_dims=config['value_hidden_dims'],
latent_dim=config['latent_dim'],
layer_norm=config['layer_norm'],
ensemble=True,
value_exp=False,
state_encoder=encoders.get('critic_state'),
goal_encoder=encoders.get('critic_goal'),
action_dim=action_dim,
)
else:
critic_def = GCBilinearValue(
hidden_dims=config['value_hidden_dims'],
latent_dim=config['latent_dim'],
layer_norm=config['layer_norm'],
ensemble=True,
value_exp=False,
state_encoder=encoders.get('critic_state'),
goal_encoder=encoders.get('critic_goal'),
)
if config['actor_loss'] == 'awr':
# AWR requires a separate V network to compute advantages (Q - V).
value_def = GCBilinearValue(
hidden_dims=config['value_hidden_dims'],
latent_dim=config['latent_dim'],
layer_norm=config['layer_norm'],
ensemble=False,
value_exp=False,
state_encoder=encoders.get('value_state'),
goal_encoder=encoders.get('value_goal'),
)
if config['discrete']:
actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=encoders.get('actor'),
)
else:
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=encoders.get('actor'),
)
network_info = dict(
critic=(critic_def, (ex_observations, ex_goals, ex_actions)),
actor=(actor_def, (ex_observations, ex_goals)),
)
if config['actor_loss'] == 'awr':
network_info.update(
value=(value_def, (ex_observations, ex_goals)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='crl', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
value_hidden_dims=(512, 512, 512), # Value network hidden dimensions.
latent_dim=512, # Latent dimension for phi and psi.
layer_norm=True, # Whether to use layer normalization.
discount=0.99, # Discount factor.
actor_loss='ddpgbc', # Actor loss type ('awr' or 'ddpgbc').
alpha=0.1, # Temperature in AWR or BC coefficient in DDPG+BC.
const_std=True, # Whether to use constant standard deviation for the actor.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='GCDataset', # Dataset class name.
value_p_curgoal=0.0, # Probability of using the current state as the value goal.
value_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the value goal.
value_p_randomgoal=0.0, # Probability of using a random state as the value goal.
value_geom_sample=True, # Whether to use geometric sampling for future value goals.
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=False, # Unused (defined for compatibility with GCDataset).
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/gcbc.py
================================================
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import GCActor, GCDiscreteActor
class GCBCAgent(flax.struct.PyTreeNode):
"""Goal-conditioned behavioral cloning (GCBC) agent."""
rng: Any
network: Any
config: Any = nonpytree_field()
def actor_loss(self, batch, grad_params, rng=None):
"""Compute the BC actor loss."""
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -log_prob.mean()
actor_info = {
'actor_loss': actor_loss,
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
rng, actor_rng = jax.random.split(rng)
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = actor_loss
return loss, info
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define encoder.
encoders = dict()
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
encoders['actor'] = GCEncoder(concat_encoder=encoder_module())
# Define actor network.
if config['discrete']:
actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=encoders.get('actor'),
)
else:
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=encoders.get('actor'),
)
network_info = dict(
actor=(actor_def, (ex_observations, ex_goals)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='gcbc', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
discount=0.99, # Discount factor (unused by default; can be used for geometric goal sampling in GCDataset).
const_std=True, # Whether to use constant standard deviation for the actor.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='GCDataset', # Dataset class name.
value_p_curgoal=0.0, # Unused (defined for compatibility with GCDataset).
value_p_trajgoal=1.0, # Unused (defined for compatibility with GCDataset).
value_p_randomgoal=0.0, # Unused (defined for compatibility with GCDataset).
value_geom_sample=False, # Unused (defined for compatibility with GCDataset).
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=True, # Unused (defined for compatibility with GCDataset).
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/gciql.py
================================================
import copy
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import GCActor, GCDiscreteActor, GCDiscreteCritic, GCValue
class GCIQLAgent(flax.struct.PyTreeNode):
"""Goal-conditioned implicit Q-learning (GCIQL) agent.
This implementation supports both AWR (actor_loss='awr') and DDPG+BC (actor_loss='ddpgbc') for the actor loss.
"""
rng: Any
network: Any
config: Any = nonpytree_field()
@staticmethod
def expectile_loss(adv, diff, expectile):
"""Compute the expectile loss."""
weight = jnp.where(adv >= 0, expectile, (1 - expectile))
return weight * (diff**2)
def value_loss(self, batch, grad_params):
"""Compute the IQL value loss."""
q1, q2 = self.network.select('target_critic')(batch['observations'], batch['value_goals'], batch['actions'])
q = jnp.minimum(q1, q2)
v = self.network.select('value')(batch['observations'], batch['value_goals'], params=grad_params)
value_loss = self.expectile_loss(q - v, q - v, self.config['expectile']).mean()
return value_loss, {
'value_loss': value_loss,
'v_mean': v.mean(),
'v_max': v.max(),
'v_min': v.min(),
}
def critic_loss(self, batch, grad_params):
"""Compute the IQL critic loss."""
next_v = self.network.select('value')(batch['next_observations'], batch['value_goals'])
q = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v
q1, q2 = self.network.select('critic')(
batch['observations'], batch['value_goals'], batch['actions'], params=grad_params
)
critic_loss = ((q1 - q) ** 2 + (q2 - q) ** 2).mean()
return critic_loss, {
'critic_loss': critic_loss,
'q_mean': q.mean(),
'q_max': q.max(),
'q_min': q.min(),
}
def actor_loss(self, batch, grad_params, rng=None):
"""Compute the actor loss (AWR or DDPG+BC)."""
if self.config['actor_loss'] == 'awr':
# AWR loss.
v = self.network.select('value')(batch['observations'], batch['actor_goals'])
q1, q2 = self.network.select('critic')(batch['observations'], batch['actor_goals'], batch['actions'])
q = jnp.minimum(q1, q2)
adv = q - v
exp_a = jnp.exp(adv * self.config['alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -(exp_a * log_prob).mean()
actor_info = {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
elif self.config['actor_loss'] == 'ddpgbc':
# DDPG+BC loss.
assert not self.config['discrete']
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
if self.config['const_std']:
q_actions = jnp.clip(dist.mode(), -1, 1)
else:
q_actions = jnp.clip(dist.sample(seed=rng), -1, 1)
q1, q2 = self.network.select('critic')(batch['observations'], batch['actor_goals'], q_actions)
q = jnp.minimum(q1, q2)
# Normalize Q values by the absolute mean to make the loss scale invariant.
q_loss = -q.mean() / jax.lax.stop_gradient(jnp.abs(q).mean() + 1e-6)
log_prob = dist.log_prob(batch['actions'])
bc_loss = -(self.config['alpha'] * log_prob).mean()
actor_loss = q_loss + bc_loss
return actor_loss, {
'actor_loss': actor_loss,
'q_loss': q_loss,
'bc_loss': bc_loss,
'q_mean': q.mean(),
'q_abs_mean': jnp.abs(q).mean(),
'bc_log_prob': log_prob.mean(),
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
else:
raise ValueError(f'Unsupported actor loss: {self.config["actor_loss"]}')
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
value_loss, value_info = self.value_loss(batch, grad_params)
for k, v in value_info.items():
info[f'value/{k}'] = v
critic_loss, critic_info = self.critic_loss(batch, grad_params)
for k, v in critic_info.items():
info[f'critic/{k}'] = v
rng, actor_rng = jax.random.split(rng)
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = value_loss + critic_loss + actor_loss
return loss, info
def target_update(self, network, module_name):
"""Update the target network."""
new_target_params = jax.tree_util.tree_map(
lambda p, tp: p * self.config['tau'] + tp * (1 - self.config['tau']),
self.network.params[f'modules_{module_name}'],
self.network.params[f'modules_target_{module_name}'],
)
network.params[f'modules_target_{module_name}'] = new_target_params
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
self.target_update(new_network, 'critic')
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define encoders.
encoders = dict()
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
encoders['value'] = GCEncoder(concat_encoder=encoder_module())
encoders['critic'] = GCEncoder(concat_encoder=encoder_module())
encoders['actor'] = GCEncoder(concat_encoder=encoder_module())
# Define value and actor networks.
value_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=False,
gc_encoder=encoders.get('value'),
)
if config['discrete']:
critic_def = GCDiscreteCritic(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
gc_encoder=encoders.get('critic'),
action_dim=action_dim,
)
else:
critic_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
gc_encoder=encoders.get('critic'),
)
if config['discrete']:
actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=encoders.get('actor'),
)
else:
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=encoders.get('actor'),
)
network_info = dict(
value=(value_def, (ex_observations, ex_goals)),
critic=(critic_def, (ex_observations, ex_goals, ex_actions)),
target_critic=(copy.deepcopy(critic_def), (ex_observations, ex_goals, ex_actions)),
actor=(actor_def, (ex_observations, ex_goals)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
params = network_params
params['modules_target_critic'] = params['modules_critic']
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='gciql', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
value_hidden_dims=(512, 512, 512), # Value network hidden dimensions.
layer_norm=True, # Whether to use layer normalization.
discount=0.99, # Discount factor.
tau=0.005, # Target network update rate.
expectile=0.9, # IQL expectile.
actor_loss='ddpgbc', # Actor loss type ('awr' or 'ddpgbc').
alpha=0.3, # Temperature in AWR or BC coefficient in DDPG+BC.
const_std=True, # Whether to use constant standard deviation for the actor.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='GCDataset', # Dataset class name.
value_p_curgoal=0.2, # Probability of using the current state as the value goal.
value_p_trajgoal=0.5, # Probability of using a future state in the same trajectory as the value goal.
value_p_randomgoal=0.3, # Probability of using a random state as the value goal.
value_geom_sample=True, # Whether to use geometric sampling for future value goals.
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=True, # Whether to use '0 if s == g else -1' (True) or '1 if s == g else 0' (False) as reward.
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/gcivl.py
================================================
import copy
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import GCActor, GCDiscreteActor, GCValue
class GCIVLAgent(flax.struct.PyTreeNode):
"""Goal-conditioned implicit V-learning (GCIVL) agent.
This is a variant of GCIQL that only uses a V function, without Q functions.
"""
rng: Any
network: Any
config: Any = nonpytree_field()
@staticmethod
def expectile_loss(adv, diff, expectile):
"""Compute the expectile loss."""
weight = jnp.where(adv >= 0, expectile, (1 - expectile))
return weight * (diff**2)
def value_loss(self, batch, grad_params):
"""Compute the IVL value loss.
This value loss is similar to the original IQL value loss, but involves additional tricks to stabilize training.
For example, when computing the expectile loss, we separate the advantage part (which is used to compute the
weight) and the difference part (which is used to compute the loss), where we use the target value function to
compute the former and the current value function to compute the latter. This is similar to how double DQN
mitigates overestimation bias.
"""
(next_v1_t, next_v2_t) = self.network.select('target_value')(batch['next_observations'], batch['value_goals'])
next_v_t = jnp.minimum(next_v1_t, next_v2_t)
q = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v_t
(v1_t, v2_t) = self.network.select('target_value')(batch['observations'], batch['value_goals'])
v_t = (v1_t + v2_t) / 2
adv = q - v_t
q1 = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v1_t
q2 = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v2_t
(v1, v2) = self.network.select('value')(batch['observations'], batch['value_goals'], params=grad_params)
v = (v1 + v2) / 2
value_loss1 = self.expectile_loss(adv, q1 - v1, self.config['expectile']).mean()
value_loss2 = self.expectile_loss(adv, q2 - v2, self.config['expectile']).mean()
value_loss = value_loss1 + value_loss2
return value_loss, {
'value_loss': value_loss,
'v_mean': v.mean(),
'v_max': v.max(),
'v_min': v.min(),
}
def actor_loss(self, batch, grad_params, rng=None):
"""Compute the AWR actor loss."""
v1, v2 = self.network.select('value')(batch['observations'], batch['actor_goals'])
nv1, nv2 = self.network.select('value')(batch['next_observations'], batch['actor_goals'])
v = (v1 + v2) / 2
nv = (nv1 + nv2) / 2
adv = nv - v
exp_a = jnp.exp(adv * self.config['alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -(exp_a * log_prob).mean()
actor_info = {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
value_loss, value_info = self.value_loss(batch, grad_params)
for k, v in value_info.items():
info[f'value/{k}'] = v
rng, actor_rng = jax.random.split(rng)
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = value_loss + actor_loss
return loss, info
def target_update(self, network, module_name):
"""Update the target network."""
new_target_params = jax.tree_util.tree_map(
lambda p, tp: p * self.config['tau'] + tp * (1 - self.config['tau']),
self.network.params[f'modules_{module_name}'],
self.network.params[f'modules_target_{module_name}'],
)
network.params[f'modules_target_{module_name}'] = new_target_params
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
self.target_update(new_network, 'value')
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define encoders.
encoders = dict()
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
encoders['value'] = GCEncoder(concat_encoder=encoder_module())
encoders['actor'] = GCEncoder(concat_encoder=encoder_module())
# Define value and actor networks.
value_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
gc_encoder=encoders.get('value'),
)
if config['discrete']:
actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=encoders.get('actor'),
)
else:
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=encoders.get('actor'),
)
network_info = dict(
value=(value_def, (ex_observations, ex_goals)),
target_value=(copy.deepcopy(value_def), (ex_observations, ex_goals)),
actor=(actor_def, (ex_observations, ex_goals)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
params = network_params
params['modules_target_value'] = params['modules_value']
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='gcivl', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
value_hidden_dims=(512, 512, 512), # Value network hidden dimensions.
layer_norm=True, # Whether to use layer normalization.
discount=0.99, # Discount factor.
tau=0.005, # Target network update rate.
expectile=0.9, # IQL expectile.
alpha=10.0, # AWR temperature.
const_std=True, # Whether to use constant standard deviation for the actor.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='GCDataset', # Dataset class name.
value_p_curgoal=0.2, # Probability of using the current state as the value goal.
value_p_trajgoal=0.5, # Probability of using a future state in the same trajectory as the value goal.
value_p_randomgoal=0.3, # Probability of using a random state as the value goal.
value_geom_sample=True, # Whether to use geometric sampling for future value goals.
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=True, # Whether to use '0 if s == g else -1' (True) or '1 if s == g else 0' (False) as reward.
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/hiql.py
================================================
from typing import Any
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import MLP, GCActor, GCDiscreteActor, GCValue, Identity, LengthNormalize
class HIQLAgent(flax.struct.PyTreeNode):
"""Hierarchical implicit Q-learning (HIQL) agent."""
rng: Any
network: Any
config: Any = nonpytree_field()
@staticmethod
def expectile_loss(adv, diff, expectile):
"""Compute the expectile loss."""
weight = jnp.where(adv >= 0, expectile, (1 - expectile))
return weight * (diff**2)
def value_loss(self, batch, grad_params):
"""Compute the IVL value loss.
This value loss is similar to the original IQL value loss, but involves additional tricks to stabilize training.
For example, when computing the expectile loss, we separate the advantage part (which is used to compute the
weight) and the difference part (which is used to compute the loss), where we use the target value function to
compute the former and the current value function to compute the latter. This is similar to how double DQN
mitigates overestimation bias.
"""
(next_v1_t, next_v2_t) = self.network.select('target_value')(batch['next_observations'], batch['value_goals'])
next_v_t = jnp.minimum(next_v1_t, next_v2_t)
q = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v_t
(v1_t, v2_t) = self.network.select('target_value')(batch['observations'], batch['value_goals'])
v_t = (v1_t + v2_t) / 2
adv = q - v_t
q1 = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v1_t
q2 = batch['rewards'] + self.config['discount'] * batch['masks'] * next_v2_t
(v1, v2) = self.network.select('value')(batch['observations'], batch['value_goals'], params=grad_params)
v = (v1 + v2) / 2
value_loss1 = self.expectile_loss(adv, q1 - v1, self.config['expectile']).mean()
value_loss2 = self.expectile_loss(adv, q2 - v2, self.config['expectile']).mean()
value_loss = value_loss1 + value_loss2
return value_loss, {
'value_loss': value_loss,
'v_mean': v.mean(),
'v_max': v.max(),
'v_min': v.min(),
}
def low_actor_loss(self, batch, grad_params):
"""Compute the low-level actor loss."""
v1, v2 = self.network.select('value')(batch['observations'], batch['low_actor_goals'])
nv1, nv2 = self.network.select('value')(batch['next_observations'], batch['low_actor_goals'])
v = (v1 + v2) / 2
nv = (nv1 + nv2) / 2
adv = nv - v
exp_a = jnp.exp(adv * self.config['low_alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
# Compute the goal representations of the subgoals.
goal_reps = self.network.select('goal_rep')(
jnp.concatenate([batch['observations'], batch['low_actor_goals']], axis=-1),
params=grad_params,
)
if not self.config['low_actor_rep_grad']:
# Stop gradients through the goal representations.
goal_reps = jax.lax.stop_gradient(goal_reps)
dist = self.network.select('low_actor')(batch['observations'], goal_reps, goal_encoded=True, params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -(exp_a * log_prob).mean()
actor_info = {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
def high_actor_loss(self, batch, grad_params):
"""Compute the high-level actor loss."""
v1, v2 = self.network.select('value')(batch['observations'], batch['high_actor_goals'])
nv1, nv2 = self.network.select('value')(batch['high_actor_targets'], batch['high_actor_goals'])
v = (v1 + v2) / 2
nv = (nv1 + nv2) / 2
adv = nv - v
exp_a = jnp.exp(adv * self.config['high_alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
dist = self.network.select('high_actor')(batch['observations'], batch['high_actor_goals'], params=grad_params)
target = self.network.select('goal_rep')(
jnp.concatenate([batch['observations'], batch['high_actor_targets']], axis=-1)
)
log_prob = dist.log_prob(target)
actor_loss = -(exp_a * log_prob).mean()
return actor_loss, {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
'mse': jnp.mean((dist.mode() - target) ** 2),
'std': jnp.mean(dist.scale_diag),
}
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
value_loss, value_info = self.value_loss(batch, grad_params)
for k, v in value_info.items():
info[f'value/{k}'] = v
low_actor_loss, low_actor_info = self.low_actor_loss(batch, grad_params)
for k, v in low_actor_info.items():
info[f'low_actor/{k}'] = v
high_actor_loss, high_actor_info = self.high_actor_loss(batch, grad_params)
for k, v in high_actor_info.items():
info[f'high_actor/{k}'] = v
loss = value_loss + low_actor_loss + high_actor_loss
return loss, info
def target_update(self, network, module_name):
"""Update the target network."""
new_target_params = jax.tree_util.tree_map(
lambda p, tp: p * self.config['tau'] + tp * (1 - self.config['tau']),
self.network.params[f'modules_{module_name}'],
self.network.params[f'modules_target_{module_name}'],
)
network.params[f'modules_target_{module_name}'] = new_target_params
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
self.target_update(new_network, 'value')
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor.
It first queries the high-level actor to obtain subgoal representations, and then queries the low-level actor
to obtain raw actions.
"""
high_seed, low_seed = jax.random.split(seed)
high_dist = self.network.select('high_actor')(observations, goals, temperature=temperature)
goal_reps = high_dist.sample(seed=high_seed)
goal_reps = goal_reps / jnp.linalg.norm(goal_reps, axis=-1, keepdims=True) * jnp.sqrt(goal_reps.shape[-1])
low_dist = self.network.select('low_actor')(observations, goal_reps, goal_encoded=True, temperature=temperature)
actions = low_dist.sample(seed=low_seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define (state-dependent) subgoal representation phi([s; g]) that outputs a length-normalized vector.
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
goal_rep_seq = [encoder_module()]
else:
goal_rep_seq = []
goal_rep_seq.append(
MLP(
hidden_dims=(*config['value_hidden_dims'], config['rep_dim']),
activate_final=False,
layer_norm=config['layer_norm'],
)
)
goal_rep_seq.append(LengthNormalize())
goal_rep_def = nn.Sequential(goal_rep_seq)
# Define the encoders that handle the inputs to the value and actor networks.
# The subgoal representation phi([s; g]) is trained by the parameterized value function V(s, phi([s; g])).
# The high-level actor predicts the subgoal representation phi([s; w]) for subgoal w given s and g.
# The low-level actor predicts actions given the current state s and the subgoal representation phi([s; w]).
if config['encoder'] is not None:
# Pixel-based environments require visual encoders for state inputs, in addition to the pre-defined shared
# encoder for subgoal representations.
# Value: V(encoder^V(s), phi([s; g]))
value_encoder_def = GCEncoder(state_encoder=encoder_module(), concat_encoder=goal_rep_def)
target_value_encoder_def = GCEncoder(state_encoder=encoder_module(), concat_encoder=goal_rep_def)
# Low-level actor: pi^l(. | encoder^l(s), phi([s; w]))
low_actor_encoder_def = GCEncoder(state_encoder=encoder_module(), concat_encoder=goal_rep_def)
# High-level actor: pi^h(. | encoder^h([s; g]))
high_actor_encoder_def = GCEncoder(concat_encoder=encoder_module())
else:
# State-based environments only use the pre-defined shared encoder for subgoal representations.
# Value: V(s, phi([s; g]))
value_encoder_def = GCEncoder(state_encoder=Identity(), concat_encoder=goal_rep_def)
target_value_encoder_def = GCEncoder(state_encoder=Identity(), concat_encoder=goal_rep_def)
# Low-level actor: pi^l(. | s, phi([s; w]))
low_actor_encoder_def = GCEncoder(state_encoder=Identity(), concat_encoder=goal_rep_def)
# High-level actor: pi^h(. | s, g) (i.e., no encoder)
high_actor_encoder_def = None
# Define value and actor networks.
value_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
gc_encoder=value_encoder_def,
)
target_value_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
gc_encoder=target_value_encoder_def,
)
if config['discrete']:
low_actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=low_actor_encoder_def,
)
else:
low_actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=low_actor_encoder_def,
)
high_actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=config['rep_dim'],
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=high_actor_encoder_def,
)
network_info = dict(
goal_rep=(goal_rep_def, (jnp.concatenate([ex_observations, ex_goals], axis=-1))),
value=(value_def, (ex_observations, ex_goals)),
target_value=(target_value_def, (ex_observations, ex_goals)),
low_actor=(low_actor_def, (ex_observations, ex_goals)),
high_actor=(high_actor_def, (ex_observations, ex_goals)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
params = network.params
params['modules_target_value'] = params['modules_value']
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='hiql', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
value_hidden_dims=(512, 512, 512), # Value network hidden dimensions.
layer_norm=True, # Whether to use layer normalization.
discount=0.99, # Discount factor.
tau=0.005, # Target network update rate.
expectile=0.7, # IQL expectile.
low_alpha=3.0, # Low-level AWR temperature.
high_alpha=3.0, # High-level AWR temperature.
subgoal_steps=25, # Subgoal steps.
rep_dim=10, # Goal representation dimension.
low_actor_rep_grad=False, # Whether low-actor gradients flow to goal representation (use True for pixels).
const_std=True, # Whether to use constant standard deviation for the actors.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='HGCDataset', # Dataset class name.
value_p_curgoal=0.2, # Probability of using the current state as the value goal.
value_p_trajgoal=0.5, # Probability of using a future state in the same trajectory as the value goal.
value_p_randomgoal=0.3, # Probability of using a random state as the value goal.
value_geom_sample=True, # Whether to use geometric sampling for future value goals.
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=True, # Whether to use '0 if s == g else -1' (True) or '1 if s == g else 0' (False) as reward.
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/qrl.py
================================================
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import optax
from utils.encoders import GCEncoder, encoder_modules
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import MLP, GCActor, GCDiscreteActor, GCIQEValue, GCMRNValue, LogParam
class QRLAgent(flax.struct.PyTreeNode):
"""Quasimetric RL (QRL) agent.
This implementation supports the following variants:
(1) Value parameterizations: IQE (quasimetric_type='iqe') and MRN (quasimetric_type='mrn').
(2) Actor losses: AWR (actor_loss='awr') and latent dynamics-based DDPG+BC (actor_loss='ddpgbc').
QRL with AWR only fits a quasimetric value function and an actor network. QRL with DDPG+BC fits a quasimetric value
function, an actor network, and a latent dynamics model. The latent dynamics model is used to compute
reparameterized gradients for the actor loss. The original implementation of QRL uses IQE and DDPG+BC.
"""
rng: Any
network: Any
config: Any = nonpytree_field()
def value_loss(self, batch, grad_params):
"""Compute the QRL value loss."""
d_neg = self.network.select('value')(batch['observations'], batch['value_goals'], params=grad_params)
d_pos = self.network.select('value')(batch['observations'], batch['next_observations'], params=grad_params)
lam = self.network.select('lam')(params=grad_params)
# Apply loss shaping following the original implementation.
d_neg_loss = (100 * jax.nn.softplus(5 - d_neg / 100)).mean()
d_pos_loss = (jax.nn.relu(d_pos - 1) ** 2).mean()
value_loss = d_neg_loss + d_pos_loss * jax.lax.stop_gradient(lam)
lam_loss = lam * (self.config['eps'] - jax.lax.stop_gradient(d_pos_loss))
total_loss = value_loss + lam_loss
return total_loss, {
'total_loss': total_loss,
'value_loss': value_loss,
'lam_loss': lam_loss,
'd_neg_loss': d_neg_loss,
'd_neg_mean': d_neg.mean(),
'd_neg_max': d_neg.max(),
'd_neg_min': d_neg.min(),
'd_pos_loss': d_pos_loss,
'd_pos_mean': d_pos.mean(),
'd_pos_max': d_pos.max(),
'd_pos_min': d_pos.min(),
'lam': lam,
}
def dynamics_loss(self, batch, grad_params):
"""Compute the dynamics loss."""
_, ob_reps, next_ob_reps = self.network.select('value')(
batch['observations'], batch['next_observations'], info=True, params=grad_params
)
# Dynamics model predicts the delta of the next observation.
pred_next_ob_reps = ob_reps + self.network.select('dynamics')(
jnp.concatenate([ob_reps, batch['actions']], axis=-1), params=grad_params
)
dist1 = self.network.select('value')(next_ob_reps, pred_next_ob_reps, is_phi=True, params=grad_params)
dist2 = self.network.select('value')(pred_next_ob_reps, next_ob_reps, is_phi=True, params=grad_params)
dynamics_loss = (dist1 + dist2).mean() / 2
return dynamics_loss, {
'dynamics_loss': dynamics_loss,
}
def actor_loss(self, batch, grad_params, rng=None):
"""Compute the actor loss (AWR or DDPG+BC)."""
if self.config['actor_loss'] == 'awr':
# Compute AWR loss based on V(s', g) - V(s, g).
v = -self.network.select('value')(batch['observations'], batch['actor_goals'])
nv = -self.network.select('value')(batch['next_observations'], batch['actor_goals'])
adv = nv - v
exp_a = jnp.exp(adv * self.config['alpha'])
exp_a = jnp.minimum(exp_a, 100.0)
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
log_prob = dist.log_prob(batch['actions'])
actor_loss = -(exp_a * log_prob).mean()
actor_info = {
'actor_loss': actor_loss,
'adv': adv.mean(),
'bc_log_prob': log_prob.mean(),
}
if not self.config['discrete']:
actor_info.update(
{
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
)
return actor_loss, actor_info
elif self.config['actor_loss'] == 'ddpgbc':
# Compute DDPG+BC loss based on latent dynamics model.
assert not self.config['discrete']
dist = self.network.select('actor')(batch['observations'], batch['actor_goals'], params=grad_params)
if self.config['const_std']:
q_actions = jnp.clip(dist.mode(), -1, 1)
else:
q_actions = jnp.clip(dist.sample(seed=rng), -1, 1)
_, ob_reps, goal_reps = self.network.select('value')(batch['observations'], batch['actor_goals'], info=True)
pred_next_ob_reps = ob_reps + self.network.select('dynamics')(
jnp.concatenate([ob_reps, q_actions], axis=-1)
)
q = -self.network.select('value')(pred_next_ob_reps, goal_reps, is_phi=True)
# Normalize Q values by the absolute mean to make the loss scale invariant.
q_loss = -q.mean() / jax.lax.stop_gradient(jnp.abs(q).mean() + 1e-6)
log_prob = dist.log_prob(batch['actions'])
bc_loss = -(self.config['alpha'] * log_prob).mean()
actor_loss = q_loss + bc_loss
return actor_loss, {
'actor_loss': actor_loss,
'q_loss': q_loss,
'bc_loss': bc_loss,
'q_mean': q.mean(),
'q_abs_mean': jnp.abs(q).mean(),
'bc_log_prob': log_prob.mean(),
'mse': jnp.mean((dist.mode() - batch['actions']) ** 2),
'std': jnp.mean(dist.scale_diag),
}
else:
raise ValueError(f'Unsupported actor loss: {self.config["actor_loss"]}')
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
value_loss, value_info = self.value_loss(batch, grad_params)
for k, v in value_info.items():
info[f'value/{k}'] = v
if self.config['actor_loss'] == 'ddpgbc':
dynamics_loss, dynamics_info = self.dynamics_loss(batch, grad_params)
for k, v in dynamics_info.items():
info[f'dynamics/{k}'] = v
else:
dynamics_loss = 0.0
rng, actor_rng = jax.random.split(rng)
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = value_loss + dynamics_loss + actor_loss
return loss, info
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
if not self.config['discrete']:
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions. In discrete-action MDPs, this should contain the maximum action value.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
ex_goals = ex_observations
ex_latents = np.zeros((ex_observations.shape[0], config['latent_dim']), dtype=np.float32)
if config['discrete']:
action_dim = ex_actions.max() + 1
else:
action_dim = ex_actions.shape[-1]
# Define encoders.
encoders = dict()
if config['encoder'] is not None:
encoder_module = encoder_modules[config['encoder']]
encoders['value'] = encoder_module()
encoders['actor'] = GCEncoder(concat_encoder=encoder_module())
# Define value and actor networks.
if config['quasimetric_type'] == 'mrn':
value_def = GCMRNValue(
hidden_dims=config['value_hidden_dims'],
latent_dim=config['latent_dim'],
layer_norm=config['layer_norm'],
encoder=encoders.get('value'),
)
elif config['quasimetric_type'] == 'iqe':
value_def = GCIQEValue(
hidden_dims=config['value_hidden_dims'],
latent_dim=config['latent_dim'],
dim_per_component=8,
layer_norm=config['layer_norm'],
encoder=encoders.get('value'),
)
else:
raise ValueError(f'Unsupported quasimetric type: {config["quasimetric_type"]}')
if config['actor_loss'] == 'ddpgbc':
# DDPG+BC requires a latent dynamics model.
dynamics_def = MLP(
hidden_dims=(*config['value_hidden_dims'], config['latent_dim']),
layer_norm=config['layer_norm'],
)
if config['discrete']:
actor_def = GCDiscreteActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
gc_encoder=encoders.get('actor'),
)
else:
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
state_dependent_std=False,
const_std=config['const_std'],
gc_encoder=encoders.get('actor'),
)
# Define the dual lambda variable.
lam_def = LogParam()
network_info = dict(
value=(value_def, (ex_observations, ex_goals)),
actor=(actor_def, (ex_observations, ex_goals)),
lam=(lam_def, ()),
)
if config['actor_loss'] == 'ddpgbc':
network_info.update(
dynamics=(dynamics_def, np.concatenate([ex_latents, ex_actions], axis=-1)),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
# Agent hyperparameters.
agent_name='qrl', # Agent name.
lr=3e-4, # Learning rate.
batch_size=1024, # Batch size.
actor_hidden_dims=(512, 512, 512), # Actor network hidden dimensions.
value_hidden_dims=(512, 512, 512), # Value network hidden dimensions.
quasimetric_type='iqe', # Quasimetric parameterization type ('iqe' or 'mrn').
latent_dim=512, # Latent dimension for the quasimetric value function.
layer_norm=True, # Whether to use layer normalization.
discount=0.99, # Discount factor (unused by default; can be used for geometric goal sampling in GCDataset).
eps=0.05, # Margin for the dual lambda loss.
actor_loss='ddpgbc', # Actor loss type ('awr' or 'ddpgbc').
alpha=0.003, # Temperature in AWR or BC coefficient in DDPG+BC.
const_std=True, # Whether to use constant standard deviation for the actor.
discrete=False, # Whether the action space is discrete.
encoder=ml_collections.config_dict.placeholder(str), # Visual encoder name (None, 'impala_small', etc.).
# Dataset hyperparameters.
dataset_class='GCDataset', # Dataset class name.
value_p_curgoal=0.0, # Probability of using the current state as the value goal.
value_p_trajgoal=0.0, # Probability of using a future state in the same trajectory as the value goal.
value_p_randomgoal=1.0, # Probability of using a random state as the value goal.
value_geom_sample=True, # Whether to use geometric sampling for future value goals.
actor_p_curgoal=0.0, # Probability of using the current state as the actor goal.
actor_p_trajgoal=1.0, # Probability of using a future state in the same trajectory as the actor goal.
actor_p_randomgoal=0.0, # Probability of using a random state as the actor goal.
actor_geom_sample=False, # Whether to use geometric sampling for future actor goals.
gc_negative=False, # Unused (defined for compatibility with GCDataset).
p_aug=0.0, # Probability of applying image augmentation.
frame_stack=ml_collections.config_dict.placeholder(int), # Number of frames to stack.
)
)
return config
================================================
FILE: impls/agents/sac.py
================================================
import copy
from typing import Any
import flax
import jax
import jax.numpy as jnp
import ml_collections
import optax
from utils.flax_utils import ModuleDict, TrainState, nonpytree_field
from utils.networks import GCActor, GCValue, LogParam
class SACAgent(flax.struct.PyTreeNode):
"""Soft actor-critic (SAC) agent."""
rng: Any
network: Any
config: Any = nonpytree_field()
def critic_loss(self, batch, grad_params, rng):
"""Compute the SAC critic loss."""
next_dist = self.network.select('actor')(batch['next_observations'])
next_actions, next_log_probs = next_dist.sample_and_log_prob(seed=rng)
next_qs = self.network.select('target_critic')(batch['next_observations'], actions=next_actions)
if self.config['min_q']:
next_q = jnp.min(next_qs, axis=0)
else:
next_q = jnp.mean(next_qs, axis=0)
target_q = batch['rewards'] + self.config['discount'] * batch['masks'] * next_q
target_q = target_q - self.config['discount'] * batch['masks'] * next_log_probs * self.network.select('alpha')()
q = self.network.select('critic')(batch['observations'], actions=batch['actions'], params=grad_params)
critic_loss = jnp.square(q - target_q).mean()
return critic_loss, {
'critic_loss': critic_loss,
'q_mean': q.mean(),
'q_max': q.max(),
'q_min': q.min(),
}
def actor_loss(self, batch, grad_params, rng):
"""Compute the SAC actor loss."""
# Actor loss.
dist = self.network.select('actor')(batch['observations'], params=grad_params)
actions, log_probs = dist.sample_and_log_prob(seed=rng)
qs = self.network.select('critic')(batch['observations'], actions=actions)
if self.config['min_q']:
q = jnp.min(qs, axis=0)
else:
q = jnp.mean(qs, axis=0)
actor_loss = (log_probs * self.network.select('alpha')() - q).mean()
# Entropy loss.
alpha = self.network.select('alpha')(params=grad_params)
entropy = -jax.lax.stop_gradient(log_probs).mean()
alpha_loss = (alpha * (entropy - self.config['target_entropy'])).mean()
total_loss = actor_loss + alpha_loss
if self.config['tanh_squash']:
action_std = dist._distribution.stddev()
else:
action_std = dist.stddev().mean()
return total_loss, {
'total_loss': total_loss,
'actor_loss': actor_loss,
'alpha_loss': alpha_loss,
'alpha': alpha,
'entropy': -log_probs.mean(),
'std': action_std.mean(),
}
@jax.jit
def total_loss(self, batch, grad_params, rng=None):
"""Compute the total loss."""
info = {}
rng = rng if rng is not None else self.rng
rng, actor_rng, critic_rng = jax.random.split(rng, 3)
critic_loss, critic_info = self.critic_loss(batch, grad_params, critic_rng)
for k, v in critic_info.items():
info[f'critic/{k}'] = v
actor_loss, actor_info = self.actor_loss(batch, grad_params, actor_rng)
for k, v in actor_info.items():
info[f'actor/{k}'] = v
loss = critic_loss + actor_loss
return loss, info
def target_update(self, network, module_name):
"""Update the target network."""
new_target_params = jax.tree_util.tree_map(
lambda p, tp: p * self.config['tau'] + tp * (1 - self.config['tau']),
self.network.params[f'modules_{module_name}'],
self.network.params[f'modules_target_{module_name}'],
)
network.params[f'modules_target_{module_name}'] = new_target_params
@jax.jit
def update(self, batch):
"""Update the agent and return a new agent with information dictionary."""
new_rng, rng = jax.random.split(self.rng)
def loss_fn(grad_params):
return self.total_loss(batch, grad_params, rng=rng)
new_network, info = self.network.apply_loss_fn(loss_fn=loss_fn)
self.target_update(new_network, 'critic')
return self.replace(network=new_network, rng=new_rng), info
@jax.jit
def sample_actions(
self,
observations,
goals=None,
seed=None,
temperature=1.0,
):
"""Sample actions from the actor."""
dist = self.network.select('actor')(observations, goals, temperature=temperature)
actions = dist.sample(seed=seed)
actions = jnp.clip(actions, -1, 1)
return actions
@classmethod
def create(
cls,
seed,
ex_observations,
ex_actions,
config,
):
"""Create a new agent.
Args:
seed: Random seed.
ex_observations: Example batch of observations.
ex_actions: Example batch of actions.
config: Configuration dictionary.
"""
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng, 2)
action_dim = ex_actions.shape[-1]
if config['target_entropy'] is None:
config['target_entropy'] = -config['target_entropy_multiplier'] * action_dim
# Define critic and actor networks.
critic_def = GCValue(
hidden_dims=config['value_hidden_dims'],
layer_norm=config['layer_norm'],
ensemble=True,
)
actor_def = GCActor(
hidden_dims=config['actor_hidden_dims'],
action_dim=action_dim,
log_std_min=-5,
tanh_squash=config['tanh_squash'],
state_dependent_std=config['state_dependent_std'],
const_std=False,
final_fc_init_scale=config['actor_fc_scale'],
)
# Define the dual alpha variable.
alpha_def = LogParam()
network_info = dict(
critic=(critic_def, (ex_observations, None, ex_actions)),
target_critic=(copy.deepcopy(critic_def), (ex_observations, None, ex_actions)),
actor=(actor_def, (ex_observations, None)),
alpha=(alpha_def, ()),
)
networks = {k: v[0] for k, v in network_info.items()}
network_args = {k: v[1] for k, v in network_info.items()}
network_def = ModuleDict(networks)
network_tx = optax.adam(learning_rate=config['lr'])
network_params = network_def.init(init_rng, **network_args)['params']
network = TrainState.create(network_def, network_params, tx=network_tx)
params = network.params
params['modules_target_critic'] = params['modules_critic']
return cls(rng, network=network, config=flax.core.FrozenDict(**config))
def get_config():
config = ml_collections.ConfigDict(
dict(
agent_name='sac', # Agent name.
lr=1e-4, # Learning rate.
batch_size=256, # Batch size.
actor_hidden_dims=(256, 256), # Actor network hidden dimensions.
value_hidden_dims=(256, 256), # Value network hidden dimensions.
layer_norm=False, # Whether to use layer normalization.
discount=0.99, # Discount factor.
tau=0.005, # Target network update rate.
target_entropy=ml_collections.config_dict.placeholder(float), # Target entropy (None for automatic tuning).
target_entropy_multiplier=0.5, # Multiplier to dim(A) for target entropy.
tanh_squash=True, # Whether to squash actions with tanh.
state_dependent_std=True, # Whether to use state-dependent standard deviations for actor.
actor_fc_scale=0.01, # Final layer initialization scale for actor.
min_q=True, # Whether to use min Q (True) or mean Q (False).
)
)
return config
================================================
FILE: impls/hyperparameters.sh
================================================
# pointmaze-medium-navigate-v0 (GCBC)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-medium-navigate-v0 (GCIVL)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# pointmaze-medium-navigate-v0 (GCIQL)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.003
# pointmaze-medium-navigate-v0 (QRL)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.0003
# pointmaze-medium-navigate-v0 (CRL)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.03
# pointmaze-medium-navigate-v0 (HIQL)
python main.py --env_name=pointmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-large-navigate-v0 (GCBC)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# pointmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.003
# pointmaze-large-navigate-v0 (QRL)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.0003
# pointmaze-large-navigate-v0 (CRL)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.03
# pointmaze-large-navigate-v0 (HIQL)
python main.py --env_name=pointmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-giant-navigate-v0 (GCBC)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-giant-navigate-v0 (GCIVL)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.discount=0.995
# pointmaze-giant-navigate-v0 (GCIQL)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.003 --agent.discount=0.995
# pointmaze-giant-navigate-v0 (QRL)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.0003 --agent.discount=0.995
# pointmaze-giant-navigate-v0 (CRL)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.03 --agent.discount=0.995
# pointmaze-giant-navigate-v0 (HIQL)
python main.py --env_name=pointmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-teleport-navigate-v0 (GCBC)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-teleport-navigate-v0 (GCIVL)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# pointmaze-teleport-navigate-v0 (GCIQL)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.003
# pointmaze-teleport-navigate-v0 (QRL)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.0003
# pointmaze-teleport-navigate-v0 (CRL)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.03
# pointmaze-teleport-navigate-v0 (HIQL)
python main.py --env_name=pointmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-medium-stitch-v0 (GCBC)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-medium-stitch-v0 (GCIVL)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# pointmaze-medium-stitch-v0 (GCIQL)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# pointmaze-medium-stitch-v0 (QRL)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.0003
# pointmaze-medium-stitch-v0 (CRL)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.03
# pointmaze-medium-stitch-v0 (HIQL)
python main.py --env_name=pointmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-large-stitch-v0 (GCBC)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-large-stitch-v0 (GCIVL)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# pointmaze-large-stitch-v0 (GCIQL)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# pointmaze-large-stitch-v0 (QRL)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.0003
# pointmaze-large-stitch-v0 (CRL)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.03
# pointmaze-large-stitch-v0 (HIQL)
python main.py --env_name=pointmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-giant-stitch-v0 (GCBC)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-giant-stitch-v0 (GCIVL)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.discount=0.995
# pointmaze-giant-stitch-v0 (GCIQL)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.discount=0.995
# pointmaze-giant-stitch-v0 (QRL)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.0003 --agent.discount=0.995
# pointmaze-giant-stitch-v0 (CRL)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.03 --agent.discount=0.995
# pointmaze-giant-stitch-v0 (HIQL)
python main.py --env_name=pointmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# pointmaze-teleport-stitch-v0 (GCBC)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# pointmaze-teleport-stitch-v0 (GCIVL)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# pointmaze-teleport-stitch-v0 (GCIQL)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# pointmaze-teleport-stitch-v0 (QRL)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.0003
# pointmaze-teleport-stitch-v0 (CRL)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.03
# pointmaze-teleport-stitch-v0 (HIQL)
python main.py --env_name=pointmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-medium-navigate-v0 (GCBC)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-medium-navigate-v0 (GCIVL)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# antmaze-medium-navigate-v0 (GCIQL)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3
# antmaze-medium-navigate-v0 (QRL)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003
# antmaze-medium-navigate-v0 (CRL)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# antmaze-medium-navigate-v0 (HIQL)
python main.py --env_name=antmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-large-navigate-v0 (GCBC)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# antmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3
# antmaze-large-navigate-v0 (QRL)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003
# antmaze-large-navigate-v0 (CRL)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# antmaze-large-navigate-v0 (HIQL)
python main.py --env_name=antmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-giant-navigate-v0 (GCBC)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-giant-navigate-v0 (GCIVL)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.discount=0.995
# antmaze-giant-navigate-v0 (GCIQL)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3 --agent.discount=0.995
# antmaze-giant-navigate-v0 (QRL)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003 --agent.discount=0.995
# antmaze-giant-navigate-v0 (CRL)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.discount=0.995
# antmaze-giant-navigate-v0 (HIQL)
python main.py --env_name=antmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-teleport-navigate-v0 (GCBC)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-teleport-navigate-v0 (GCIVL)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# antmaze-teleport-navigate-v0 (GCIQL)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3
# antmaze-teleport-navigate-v0 (QRL)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003
# antmaze-teleport-navigate-v0 (CRL)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# antmaze-teleport-navigate-v0 (HIQL)
python main.py --env_name=antmaze-teleport-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-medium-stitch-v0 (GCBC)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-medium-stitch-v0 (GCIVL)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# antmaze-medium-stitch-v0 (GCIQL)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3
# antmaze-medium-stitch-v0 (QRL)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# antmaze-medium-stitch-v0 (CRL)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1
# antmaze-medium-stitch-v0 (HIQL)
python main.py --env_name=antmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-large-stitch-v0 (GCBC)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-large-stitch-v0 (GCIVL)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# antmaze-large-stitch-v0 (GCIQL)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3
# antmaze-large-stitch-v0 (QRL)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# antmaze-large-stitch-v0 (CRL)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1
# antmaze-large-stitch-v0 (HIQL)
python main.py --env_name=antmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-giant-stitch-v0 (GCBC)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-giant-stitch-v0 (GCIVL)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.discount=0.995
# antmaze-giant-stitch-v0 (GCIQL)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3 --agent.discount=0.995
# antmaze-giant-stitch-v0 (QRL)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.discount=0.995
# antmaze-giant-stitch-v0 (CRL)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# antmaze-giant-stitch-v0 (HIQL)
python main.py --env_name=antmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-teleport-stitch-v0 (GCBC)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-teleport-stitch-v0 (GCIVL)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# antmaze-teleport-stitch-v0 (GCIQL)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3
# antmaze-teleport-stitch-v0 (QRL)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# antmaze-teleport-stitch-v0 (CRL)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1
# antmaze-teleport-stitch-v0 (HIQL)
python main.py --env_name=antmaze-teleport-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antmaze-medium-explore-v0 (GCBC)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-medium-explore-v0 (GCIVL)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0
# antmaze-medium-explore-v0 (GCIQL)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01
# antmaze-medium-explore-v0 (QRL)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001
# antmaze-medium-explore-v0 (CRL)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003
# antmaze-medium-explore-v0 (HIQL)
python main.py --env_name=antmaze-medium-explore-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.high_alpha=10.0 --agent.low_alpha=10.0
# antmaze-large-explore-v0 (GCBC)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-large-explore-v0 (GCIVL)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0
# antmaze-large-explore-v0 (GCIQL)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01
# antmaze-large-explore-v0 (QRL)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001
# antmaze-large-explore-v0 (CRL)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003
# antmaze-large-explore-v0 (HIQL)
python main.py --env_name=antmaze-large-explore-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.high_alpha=10.0 --agent.low_alpha=10.0
# antmaze-teleport-explore-v0 (GCBC)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antmaze-teleport-explore-v0 (GCIVL)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0
# antmaze-teleport-explore-v0 (GCIQL)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01
# antmaze-teleport-explore-v0 (QRL)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001
# antmaze-teleport-explore-v0 (CRL)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003
# antmaze-teleport-explore-v0 (HIQL)
python main.py --env_name=antmaze-teleport-explore-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.high_alpha=10.0 --agent.low_alpha=10.0
# humanoidmaze-medium-navigate-v0 (GCBC)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-medium-navigate-v0 (GCIVL)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-medium-navigate-v0 (GCIQL)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-medium-navigate-v0 (QRL)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-medium-navigate-v0 (CRL)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-medium-navigate-v0 (HIQL)
python main.py --env_name=humanoidmaze-medium-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# humanoidmaze-large-navigate-v0 (GCBC)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-large-navigate-v0 (QRL)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-large-navigate-v0 (CRL)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-large-navigate-v0 (HIQL)
python main.py --env_name=humanoidmaze-large-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# humanoidmaze-giant-navigate-v0 (GCBC)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-giant-navigate-v0 (GCIVL)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-giant-navigate-v0 (GCIQL)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-giant-navigate-v0 (QRL)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-giant-navigate-v0 (CRL)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-giant-navigate-v0 (HIQL)
python main.py --env_name=humanoidmaze-giant-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# humanoidmaze-medium-stitch-v0 (GCBC)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-medium-stitch-v0 (GCIVL)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-medium-stitch-v0 (GCIQL)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-medium-stitch-v0 (QRL)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-medium-stitch-v0 (CRL)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-medium-stitch-v0 (HIQL)
python main.py --env_name=humanoidmaze-medium-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# humanoidmaze-large-stitch-v0 (GCBC)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-large-stitch-v0 (GCIVL)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-large-stitch-v0 (GCIQL)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-large-stitch-v0 (QRL)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-large-stitch-v0 (CRL)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-large-stitch-v0 (HIQL)
python main.py --env_name=humanoidmaze-large-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# humanoidmaze-giant-stitch-v0 (GCBC)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# humanoidmaze-giant-stitch-v0 (GCIVL)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.discount=0.995
# humanoidmaze-giant-stitch-v0 (GCIQL)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-giant-stitch-v0 (QRL)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.discount=0.995
# humanoidmaze-giant-stitch-v0 (CRL)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.discount=0.995
# humanoidmaze-giant-stitch-v0 (HIQL)
python main.py --env_name=humanoidmaze-giant-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.discount=0.995 --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=100
# antsoccer-arena-navigate-v0 (GCBC)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antsoccer-arena-navigate-v0 (GCIVL)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# antsoccer-arena-navigate-v0 (GCIQL)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1
# antsoccer-arena-navigate-v0 (QRL)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003
# antsoccer-arena-navigate-v0 (CRL)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.3
# antsoccer-arena-navigate-v0 (HIQL)
python main.py --env_name=antsoccer-arena-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antsoccer-medium-navigate-v0 (GCBC)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antsoccer-medium-navigate-v0 (GCIVL)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# antsoccer-medium-navigate-v0 (GCIQL)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1
# antsoccer-medium-navigate-v0 (QRL)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.003
# antsoccer-medium-navigate-v0 (CRL)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.3
# antsoccer-medium-navigate-v0 (HIQL)
python main.py --env_name=antsoccer-medium-navigate-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antsoccer-arena-stitch-v0 (GCBC)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antsoccer-arena-stitch-v0 (GCIVL)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# antsoccer-arena-stitch-v0 (GCIQL)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1
# antsoccer-arena-stitch-v0 (QRL)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# antsoccer-arena-stitch-v0 (CRL)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3
# antsoccer-arena-stitch-v0 (HIQL)
python main.py --env_name=antsoccer-arena-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# antsoccer-medium-stitch-v0 (GCBC)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcbc.py
# antsoccer-medium-stitch-v0 (GCIVL)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0
# antsoccer-medium-stitch-v0 (GCIQL)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1
# antsoccer-medium-stitch-v0 (QRL)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003
# antsoccer-medium-stitch-v0 (CRL)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3
# antsoccer-medium-stitch-v0 (HIQL)
python main.py --env_name=antsoccer-medium-stitch-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.high_alpha=3.0 --agent.low_alpha=3.0
# visual-antmaze-medium-navigate-v0 (GCBC)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-navigate-v0 (GCIVL)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-navigate-v0 (GCIQL)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-navigate-v0 (QRL)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-navigate-v0 (CRL)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-navigate-v0 (HIQL)
python main.py --env_name=visual-antmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-large-navigate-v0 (GCBC)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-navigate-v0 (QRL)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-navigate-v0 (CRL)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-navigate-v0 (HIQL)
python main.py --env_name=visual-antmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-giant-navigate-v0 (GCBC)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-giant-navigate-v0 (GCIVL)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-navigate-v0 (GCIQL)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.3 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-navigate-v0 (QRL)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.003 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-navigate-v0 (CRL)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-navigate-v0 (HIQL)
python main.py --env_name=visual-antmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-teleport-navigate-v0 (GCBC)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-navigate-v0 (GCIVL)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-navigate-v0 (GCIQL)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-navigate-v0 (QRL)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-navigate-v0 (CRL)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-navigate-v0 (HIQL)
python main.py --env_name=visual-antmaze-teleport-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-medium-stitch-v0 (GCBC)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-stitch-v0 (GCIVL)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-stitch-v0 (GCIQL)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-stitch-v0 (QRL)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-stitch-v0 (CRL)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-stitch-v0 (HIQL)
python main.py --env_name=visual-antmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-large-stitch-v0 (GCBC)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-stitch-v0 (GCIVL)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-stitch-v0 (GCIQL)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-stitch-v0 (QRL)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-stitch-v0 (CRL)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-stitch-v0 (HIQL)
python main.py --env_name=visual-antmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-giant-stitch-v0 (GCBC)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-giant-stitch-v0 (GCIVL)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-stitch-v0 (GCIQL)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-stitch-v0 (QRL)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-stitch-v0 (CRL)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-antmaze-giant-stitch-v0 (HIQL)
python main.py --env_name=visual-antmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-teleport-stitch-v0 (GCBC)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-stitch-v0 (GCIVL)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-stitch-v0 (GCIQL)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-stitch-v0 (QRL)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-stitch-v0 (CRL)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-stitch-v0 (HIQL)
python main.py --env_name=visual-antmaze-teleport-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0
# visual-antmaze-medium-explore-v0 (GCBC)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-explore-v0 (GCIVL)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-explore-v0 (GCIQL)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-explore-v0 (QRL)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-explore-v0 (CRL)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-medium-explore-v0 (HIQL)
python main.py --env_name=visual-antmaze-medium-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=10.0 --agent.low_actor_rep_grad=True --agent.low_alpha=10.0
# visual-antmaze-large-explore-v0 (GCBC)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-explore-v0 (GCIVL)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-explore-v0 (GCIQL)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-explore-v0 (QRL)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-explore-v0 (CRL)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-large-explore-v0 (HIQL)
python main.py --env_name=visual-antmaze-large-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=10.0 --agent.low_actor_rep_grad=True --agent.low_alpha=10.0
# visual-antmaze-teleport-explore-v0 (GCBC)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-explore-v0 (GCIVL)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-explore-v0 (GCIQL)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.01 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-explore-v0 (QRL)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.001 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-explore-v0 (CRL)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.alpha=0.003 --agent.batch_size=256 --agent.encoder=impala_small
# visual-antmaze-teleport-explore-v0 (HIQL)
python main.py --env_name=visual-antmaze-teleport-explore-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=1.0 --agent.actor_p_trajgoal=0.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=10.0 --agent.low_actor_rep_grad=True --agent.low_alpha=10.0
# visual-humanoidmaze-medium-navigate-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-medium-navigate-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-navigate-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-navigate-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-navigate-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-navigate-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-medium-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/hiql.py --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# visual-humanoidmaze-large-navigate-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-large-navigate-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-navigate-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-navigate-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-navigate-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-navigate-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-large-navigate-v0 --train_steps=500000 --eval_episodes=50 --agent=agents/hiql.py --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# visual-humanoidmaze-giant-navigate-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-giant-navigate-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-navigate-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-navigate-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-navigate-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-navigate-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-giant-navigate-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# visual-humanoidmaze-medium-stitch-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-medium-stitch-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-stitch-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-stitch-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-stitch-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-medium-stitch-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-medium-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# visual-humanoidmaze-large-stitch-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-large-stitch-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-stitch-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-stitch-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-stitch-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-large-stitch-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-large-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# visual-humanoidmaze-giant-stitch-v0 (GCBC)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small
# visual-humanoidmaze-giant-stitch-v0 (GCIVL)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=10.0 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-stitch-v0 (GCIQL)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-stitch-v0 (QRL)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.001 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-stitch-v0 (CRL)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.alpha=0.1 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small
# visual-humanoidmaze-giant-stitch-v0 (HIQL)
python main.py --env_name=visual-humanoidmaze-giant-stitch-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.actor_p_randomgoal=0.5 --agent.actor_p_trajgoal=0.5 --agent.batch_size=256 --agent.discount=0.995 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=100
# cube-single-play-v0 (GCBC)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-single-play-v0 (GCIVL)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-single-play-v0 (GCIQL)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# cube-single-play-v0 (QRL)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# cube-single-play-v0 (CRL)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# cube-single-play-v0 (HIQL)
python main.py --env_name=cube-single-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-double-play-v0 (GCBC)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-double-play-v0 (GCIVL)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-double-play-v0 (GCIQL)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# cube-double-play-v0 (QRL)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# cube-double-play-v0 (CRL)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# cube-double-play-v0 (HIQL)
python main.py --env_name=cube-double-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-triple-play-v0 (GCBC)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-triple-play-v0 (GCIVL)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-triple-play-v0 (GCIQL)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# cube-triple-play-v0 (QRL)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# cube-triple-play-v0 (CRL)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# cube-triple-play-v0 (HIQL)
python main.py --env_name=cube-triple-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-quadruple-play-v0 (GCBC)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-quadruple-play-v0 (GCIVL)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-quadruple-play-v0 (GCIQL)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# cube-quadruple-play-v0 (QRL)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# cube-quadruple-play-v0 (CRL)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# cube-quadruple-play-v0 (HIQL)
python main.py --env_name=cube-quadruple-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-single-noisy-v0 (GCBC)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-single-noisy-v0 (GCIVL)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-single-noisy-v0 (GCIQL)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# cube-single-noisy-v0 (QRL)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# cube-single-noisy-v0 (CRL)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# cube-single-noisy-v0 (HIQL)
python main.py --env_name=cube-single-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-double-noisy-v0 (GCBC)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-double-noisy-v0 (GCIVL)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-double-noisy-v0 (GCIQL)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# cube-double-noisy-v0 (QRL)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# cube-double-noisy-v0 (CRL)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# cube-double-noisy-v0 (HIQL)
python main.py --env_name=cube-double-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-triple-noisy-v0 (GCBC)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-triple-noisy-v0 (GCIVL)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-triple-noisy-v0 (GCIQL)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# cube-triple-noisy-v0 (QRL)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# cube-triple-noisy-v0 (CRL)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# cube-triple-noisy-v0 (HIQL)
python main.py --env_name=cube-triple-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# cube-quadruple-noisy-v0 (GCBC)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# cube-quadruple-noisy-v0 (GCIVL)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# cube-quadruple-noisy-v0 (GCIQL)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# cube-quadruple-noisy-v0 (QRL)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# cube-quadruple-noisy-v0 (CRL)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# cube-quadruple-noisy-v0 (HIQL)
python main.py --env_name=cube-quadruple-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# scene-play-v0 (GCBC)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# scene-play-v0 (GCIVL)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# scene-play-v0 (GCIQL)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# scene-play-v0 (QRL)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# scene-play-v0 (CRL)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# scene-play-v0 (HIQL)
python main.py --env_name=scene-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# scene-noisy-v0 (GCBC)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# scene-noisy-v0 (GCIVL)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# scene-noisy-v0 (GCIQL)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# scene-noisy-v0 (QRL)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# scene-noisy-v0 (CRL)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# scene-noisy-v0 (HIQL)
python main.py --env_name=scene-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-3x3-play-v0 (GCBC)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-3x3-play-v0 (GCIVL)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-3x3-play-v0 (GCIQL)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# puzzle-3x3-play-v0 (QRL)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# puzzle-3x3-play-v0 (CRL)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# puzzle-3x3-play-v0 (HIQL)
python main.py --env_name=puzzle-3x3-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x4-play-v0 (GCBC)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x4-play-v0 (GCIVL)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x4-play-v0 (GCIQL)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# puzzle-4x4-play-v0 (QRL)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# puzzle-4x4-play-v0 (CRL)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# puzzle-4x4-play-v0 (HIQL)
python main.py --env_name=puzzle-4x4-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x5-play-v0 (GCBC)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x5-play-v0 (GCIVL)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x5-play-v0 (GCIQL)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# puzzle-4x5-play-v0 (QRL)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# puzzle-4x5-play-v0 (CRL)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# puzzle-4x5-play-v0 (HIQL)
python main.py --env_name=puzzle-4x5-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x6-play-v0 (GCBC)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x6-play-v0 (GCIVL)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x6-play-v0 (GCIQL)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=1.0
# puzzle-4x6-play-v0 (QRL)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.3
# puzzle-4x6-play-v0 (CRL)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=3.0
# puzzle-4x6-play-v0 (HIQL)
python main.py --env_name=puzzle-4x6-play-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-3x3-noisy-v0 (GCBC)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-3x3-noisy-v0 (GCIVL)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-3x3-noisy-v0 (GCIQL)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# puzzle-3x3-noisy-v0 (QRL)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# puzzle-3x3-noisy-v0 (CRL)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# puzzle-3x3-noisy-v0 (HIQL)
python main.py --env_name=puzzle-3x3-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x4-noisy-v0 (GCBC)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x4-noisy-v0 (GCIVL)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x4-noisy-v0 (GCIQL)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# puzzle-4x4-noisy-v0 (QRL)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# puzzle-4x4-noisy-v0 (CRL)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# puzzle-4x4-noisy-v0 (HIQL)
python main.py --env_name=puzzle-4x4-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x5-noisy-v0 (GCBC)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x5-noisy-v0 (GCIVL)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x5-noisy-v0 (GCIQL)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# puzzle-4x5-noisy-v0 (QRL)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# puzzle-4x5-noisy-v0 (CRL)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# puzzle-4x5-noisy-v0 (HIQL)
python main.py --env_name=puzzle-4x5-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# puzzle-4x6-noisy-v0 (GCBC)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/gcbc.py
# puzzle-4x6-noisy-v0 (GCIVL)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/gcivl.py --agent.alpha=10.0
# puzzle-4x6-noisy-v0 (GCIQL)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/gciql.py --agent.alpha=0.03
# puzzle-4x6-noisy-v0 (QRL)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/qrl.py --agent.alpha=0.03
# puzzle-4x6-noisy-v0 (CRL)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/crl.py --agent.alpha=0.1
# puzzle-4x6-noisy-v0 (HIQL)
python main.py --env_name=puzzle-4x6-noisy-v0 --eval_episodes=50 --agent=agents/hiql.py --agent.high_alpha=3.0 --agent.low_alpha=3.0 --agent.subgoal_steps=10
# visual-cube-single-play-v0 (GCBC)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-play-v0 (GCIVL)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-play-v0 (GCIQL)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-play-v0 (QRL)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-play-v0 (CRL)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-play-v0 (HIQL)
python main.py --env_name=visual-cube-single-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-double-play-v0 (GCBC)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-play-v0 (GCIVL)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-play-v0 (GCIQL)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-play-v0 (QRL)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-play-v0 (CRL)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-play-v0 (HIQL)
python main.py --env_name=visual-cube-double-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-triple-play-v0 (GCBC)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-play-v0 (GCIVL)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-play-v0 (GCIQL)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-play-v0 (QRL)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-play-v0 (CRL)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-play-v0 (HIQL)
python main.py --env_name=visual-cube-triple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-quadruple-play-v0 (GCBC)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-play-v0 (GCIVL)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-play-v0 (GCIQL)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-play-v0 (QRL)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-play-v0 (CRL)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-play-v0 (HIQL)
python main.py --env_name=visual-cube-quadruple-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-single-noisy-v0 (GCBC)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-noisy-v0 (GCIVL)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-noisy-v0 (GCIQL)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-noisy-v0 (QRL)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-noisy-v0 (CRL)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-single-noisy-v0 (HIQL)
python main.py --env_name=visual-cube-single-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-double-noisy-v0 (GCBC)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-noisy-v0 (GCIVL)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-noisy-v0 (GCIQL)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-noisy-v0 (QRL)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-noisy-v0 (CRL)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-double-noisy-v0 (HIQL)
python main.py --env_name=visual-cube-double-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-triple-noisy-v0 (GCBC)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-noisy-v0 (GCIVL)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-noisy-v0 (GCIQL)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-noisy-v0 (QRL)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-noisy-v0 (CRL)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-triple-noisy-v0 (HIQL)
python main.py --env_name=visual-cube-triple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-cube-quadruple-noisy-v0 (GCBC)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-noisy-v0 (GCIVL)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-noisy-v0 (GCIQL)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-noisy-v0 (QRL)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-noisy-v0 (CRL)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-cube-quadruple-noisy-v0 (HIQL)
python main.py --env_name=visual-cube-quadruple-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-scene-play-v0 (GCBC)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-play-v0 (GCIVL)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-play-v0 (GCIQL)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-play-v0 (QRL)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-play-v0 (CRL)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-play-v0 (HIQL)
python main.py --env_name=visual-scene-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-scene-noisy-v0 (GCBC)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-noisy-v0 (GCIVL)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-noisy-v0 (GCIQL)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-noisy-v0 (QRL)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-noisy-v0 (CRL)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-scene-noisy-v0 (HIQL)
python main.py --env_name=visual-scene-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-3x3-play-v0 (GCBC)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-play-v0 (GCIVL)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-play-v0 (GCIQL)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-play-v0 (QRL)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-play-v0 (CRL)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-play-v0 (HIQL)
python main.py --env_name=visual-puzzle-3x3-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x4-play-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-play-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-play-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-play-v0 (QRL)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-play-v0 (CRL)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-play-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x4-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x5-play-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-play-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-play-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-play-v0 (QRL)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-play-v0 (CRL)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-play-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x5-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x6-play-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-play-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-play-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=1.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-play-v0 (QRL)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.3 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-play-v0 (CRL)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-play-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x6-play-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-3x3-noisy-v0 (GCBC)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-noisy-v0 (GCIVL)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-noisy-v0 (GCIQL)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-noisy-v0 (QRL)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-noisy-v0 (CRL)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-3x3-noisy-v0 (HIQL)
python main.py --env_name=visual-puzzle-3x3-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x4-noisy-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-noisy-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-noisy-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-noisy-v0 (QRL)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-noisy-v0 (CRL)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x4-noisy-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x4-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x5-noisy-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-noisy-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-noisy-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-noisy-v0 (QRL)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-noisy-v0 (CRL)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x5-noisy-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x5-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# visual-puzzle-4x6-noisy-v0 (GCBC)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-noisy-v0 (GCIVL)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=10.0 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-noisy-v0 (GCIQL)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/gciql.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-noisy-v0 (QRL)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/qrl.py --agent.alpha=0.03 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-noisy-v0 (CRL)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/crl.py --agent.alpha=0.1 --agent.batch_size=256 --agent.encoder=impala_small --agent.p_aug=0.5
# visual-puzzle-4x6-noisy-v0 (HIQL)
python main.py --env_name=visual-puzzle-4x6-noisy-v0 --train_steps=500000 --eval_episodes=50 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.p_aug=0.5 --agent.subgoal_steps=10
# powderworld-easy-play-v0 (GCBC)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-easy-play-v0 (GCIVL)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-easy-play-v0 (GCIQL)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-easy-play-v0 (QRL)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-easy-play-v0 (CRL)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-easy-play-v0 (HIQL)
python main.py --env_name=powderworld-easy-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=10
# powderworld-medium-play-v0 (GCBC)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-medium-play-v0 (GCIVL)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-medium-play-v0 (GCIQL)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-medium-play-v0 (QRL)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-medium-play-v0 (CRL)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-medium-play-v0 (HIQL)
python main.py --env_name=powderworld-medium-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=10
# powderworld-hard-play-v0 (GCBC)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcbc.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-hard-play-v0 (GCIVL)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gcivl.py --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-hard-play-v0 (GCIQL)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/gciql.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-hard-play-v0 (QRL)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/qrl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-hard-play-v0 (CRL)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/crl.py --agent.actor_loss=awr --agent.alpha=3.0 --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small
# powderworld-hard-play-v0 (HIQL)
python main.py --env_name=powderworld-hard-play-v0 --train_steps=500000 --eval_episodes=50 --eval_temperature=0.3 --eval_on_cpu=0 --agent=agents/hiql.py --agent.batch_size=256 --agent.discrete=True --agent.encoder=impala_small --agent.high_alpha=3.0 --agent.low_actor_rep_grad=True --agent.low_alpha=3.0 --agent.subgoal_steps=10
================================================
FILE: impls/main.py
================================================
import json
import os
import random
import time
from collections import defaultdict
import jax
import numpy as np
import tqdm
import wandb
from absl import app, flags
from agents import agents
from ml_collections import config_flags
from utils.datasets import Dataset, GCDataset, HGCDataset
from utils.env_utils import make_env_and_datasets
from utils.evaluation import evaluate
from utils.flax_utils import restore_agent, save_agent
from utils.log_utils import CsvLogger, get_exp_name, get_flag_dict, get_wandb_video, setup_wandb
FLAGS = flags.FLAGS
flags.DEFINE_string('run_group', 'Debug', 'Run group.')
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'antmaze-large-navigate-v0', 'Environment (dataset) name.')
flags.DEFINE_string('save_dir', 'exp/', 'Save directory.')
flags.DEFINE_string('restore_path', None, 'Restore path.')
flags.DEFINE_integer('restore_epoch', None, 'Restore epoch.')
flags.DEFINE_integer('train_steps', 1000000, 'Number of training steps.')
flags.DEFINE_integer('log_interval', 5000, 'Logging interval.')
flags.DEFINE_integer('eval_interval', 100000, 'Evaluation interval.')
flags.DEFINE_integer('save_interval', 1000000, 'Saving interval.')
flags.DEFINE_integer('eval_tasks', None, 'Number of tasks to evaluate (None for all).')
flags.DEFINE_integer('eval_episodes', 20, 'Number of episodes for each task.')
flags.DEFINE_float('eval_temperature', 0, 'Actor temperature for evaluation.')
flags.DEFINE_float('eval_gaussian', None, 'Action Gaussian noise for evaluation.')
flags.DEFINE_integer('video_episodes', 1, 'Number of video episodes for each task.')
flags.DEFINE_integer('video_frame_skip', 3, 'Frame skip for videos.')
flags.DEFINE_integer('eval_on_cpu', 1, 'Whether to evaluate on CPU.')
config_flags.DEFINE_config_file('agent', 'agents/gciql.py', lock_config=False)
def main(_):
# Set up logger.
exp_name = get_exp_name(FLAGS.seed)
setup_wandb(project='OGBench', group=FLAGS.run_group, name=exp_name)
FLAGS.save_dir = os.path.join(FLAGS.save_dir, wandb.run.project, FLAGS.run_group, exp_name)
os.makedirs(FLAGS.save_dir, exist_ok=True)
flag_dict = get_flag_dict()
with open(os.path.join(FLAGS.save_dir, 'flags.json'), 'w') as f:
json.dump(flag_dict, f)
# Set up environment and dataset.
config = FLAGS.agent
env, train_dataset, val_dataset = make_env_and_datasets(FLAGS.env_name, frame_stack=config['frame_stack'])
dataset_class = {
'GCDataset': GCDataset,
'HGCDataset': HGCDataset,
}[config['dataset_class']]
train_dataset = dataset_class(Dataset.create(**train_dataset), config)
if val_dataset is not None:
val_dataset = dataset_class(Dataset.create(**val_dataset), config)
# Initialize agent.
random.seed(FLAGS.seed)
np.random.seed(FLAGS.seed)
example_batch = train_dataset.sample(1)
if config['discrete']:
# Fill with the maximum action to let the agent know the action space size.
example_batch['actions'] = np.full_like(example_batch['actions'], env.action_space.n - 1)
agent_class = agents[config['agent_name']]
agent = agent_class.create(
FLAGS.seed,
example_batch['observations'],
example_batch['actions'],
config,
)
# Restore agent.
if FLAGS.restore_path is not None:
agent = restore_agent(agent, FLAGS.restore_path, FLAGS.restore_epoch)
# Train agent.
train_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'train.csv'))
eval_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'eval.csv'))
first_time = time.time()
last_time = time.time()
for i in tqdm.tqdm(range(1, FLAGS.train_steps + 1), smoothing=0.1, dynamic_ncols=True):
# Update agent.
batch = train_dataset.sample(config['batch_size'])
agent, update_info = agent.update(batch)
# Log metrics.
if i % FLAGS.log_interval == 0:
train_metrics = {f'training/{k}': v for k, v in update_info.items()}
if val_dataset is not None:
val_batch = val_dataset.sample(config['batch_size'])
_, val_info = agent.total_loss(val_batch, grad_params=None)
train_metrics.update({f'validation/{k}': v for k, v in val_info.items()})
train_metrics['time/epoch_time'] = (time.time() - last_time) / FLAGS.log_interval
train_metrics['time/total_time'] = time.time() - first_time
last_time = time.time()
wandb.log(train_metrics, step=i)
train_logger.log(train_metrics, step=i)
# Evaluate agent.
if i == 1 or i % FLAGS.eval_interval == 0:
if FLAGS.eval_on_cpu:
eval_agent = jax.device_put(agent, device=jax.devices('cpu')[0])
else:
eval_agent = agent
renders = []
eval_metrics = {}
overall_metrics = defaultdict(list)
task_infos = env.unwrapped.task_infos if hasattr(env.unwrapped, 'task_infos') else env.task_infos
num_tasks = FLAGS.eval_tasks if FLAGS.eval_tasks is not None else len(task_infos)
for task_id in tqdm.trange(1, num_tasks + 1):
task_name = task_infos[task_id - 1]['task_name']
eval_info, trajs, cur_renders = evaluate(
agent=eval_agent,
env=env,
task_id=task_id,
config=config,
num_eval_episodes=FLAGS.eval_episodes,
num_video_episodes=FLAGS.video_episodes,
video_frame_skip=FLAGS.video_frame_skip,
eval_temperature=FLAGS.eval_temperature,
eval_gaussian=FLAGS.eval_gaussian,
)
renders.extend(cur_renders)
metric_names = ['success']
eval_metrics.update(
{f'evaluation/{task_name}_{k}': v for k, v in eval_info.items() if k in metric_names}
)
for k, v in eval_info.items():
if k in metric_names:
overall_metrics[k].append(v)
for k, v in overall_metrics.items():
eval_metrics[f'evaluation/overall_{k}'] = np.mean(v)
if FLAGS.video_episodes > 0:
video = get_wandb_video(renders=renders, n_cols=num_tasks)
eval_metrics['video'] = video
wandb.log(eval_metrics, step=i)
eval_logger.log(eval_metrics, step=i)
# Save agent.
if i % FLAGS.save_interval == 0:
save_agent(agent, FLAGS.save_dir, i)
train_logger.close()
eval_logger.close()
if __name__ == '__main__':
app.run(main)
================================================
FILE: impls/requirements.txt
================================================
ogbench # Use the PyPI version of OGBench. Replace this with `pip install -e .` if you want to use the local version.
jax[cuda12] >= 0.4.26
flax >= 0.8.4
distrax >= 0.1.5
ml_collections
matplotlib
moviepy
wandb
================================================
FILE: impls/utils/__init__.py
================================================
================================================
FILE: impls/utils/datasets.py
================================================
import dataclasses
from functools import partial
from typing import Any
import jax
import jax.numpy as jnp
import numpy as np
from flax.core.frozen_dict import FrozenDict
def get_size(data):
"""Return the size of the dataset."""
sizes = jax.tree_util.tree_map(lambda arr: len(arr), data)
return max(jax.tree_util.tree_leaves(sizes))
@partial(jax.jit, static_argnames=('padding',))
def random_crop(img, crop_from, padding):
"""Randomly crop an image.
Args:
img: Image to crop.
crop_from: Coordinates to crop from.
padding: Padding size.
"""
padded_img = jnp.pad(img, ((padding, padding), (padding, padding), (0, 0)), mode='edge')
return jax.lax.dynamic_slice(padded_img, crop_from, img.shape)
@partial(jax.jit, static_argnames=('padding',))
def batched_random_crop(imgs, crop_froms, padding):
"""Batched version of random_crop."""
return jax.vmap(random_crop, (0, 0, None))(imgs, crop_froms, padding)
class Dataset(FrozenDict):
"""Dataset class.
This class supports both regular datasets (i.e., storing both observations and next_observations) and
compact datasets (i.e., storing only observations). It assumes 'observations' is always present in the keys. If
'next_observations' is not present, it will be inferred from 'observations' by shifting the indices by 1. In this
case, set 'valids' appropriately to mask out the last state of each trajectory.
"""
@classmethod
def create(cls, freeze=True, **fields):
"""Create a dataset from the fields.
Args:
freeze: Whether to freeze the arrays.
**fields: Keys and values of the dataset.
"""
data = fields
assert 'observations' in data
if freeze:
jax.tree_util.tree_map(lambda arr: arr.setflags(write=False), data)
return cls(data)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.size = get_size(self._dict)
if 'valids' in self._dict:
(self.valid_idxs,) = np.nonzero(self['valids'] > 0)
def get_random_idxs(self, num_idxs):
"""Return `num_idxs` random indices."""
if 'valids' in self._dict:
return self.valid_idxs[np.random.randint(len(self.valid_idxs), size=num_idxs)]
else:
return np.random.randint(self.size, size=num_idxs)
def sample(self, batch_size, idxs=None):
"""Sample a batch of transitions."""
if idxs is None:
idxs = self.get_random_idxs(batch_size)
return self.get_subset(idxs)
def get_subset(self, idxs):
"""Return a subset of the dataset given the indices."""
result = jax.tree_util.tree_map(lambda arr: arr[idxs], self._dict)
if 'next_observations' not in result:
result['next_observations'] = self._dict['observations'][np.minimum(idxs + 1, self.size - 1)]
return result
class ReplayBuffer(Dataset):
"""Replay buffer class.
This class extends Dataset to support adding transitions.
"""
@classmethod
def create(cls, transition, size):
"""Create a replay buffer from the example transition.
Args:
transition: Example transition (dict).
size: Size of the replay buffer.
"""
def create_buffer(example):
example = np.array(example)
return np.zeros((size, *example.shape), dtype=example.dtype)
buffer_dict = jax.tree_util.tree_map(create_buffer, transition)
return cls(buffer_dict)
@classmethod
def create_from_initial_dataset(cls, init_dataset, size):
"""Create a replay buffer from the initial dataset.
Args:
init_dataset: Initial dataset.
size: Size of the replay buffer.
"""
def create_buffer(init_buffer):
buffer = np.zeros((size, *init_buffer.shape[1:]), dtype=init_buffer.dtype)
buffer[: len(init_buffer)] = init_buffer
return buffer
buffer_dict = jax.tree_util.tree_map(create_buffer, init_dataset)
dataset = cls(buffer_dict)
dataset.size = dataset.pointer = get_size(init_dataset)
return dataset
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.max_size = get_size(self._dict)
self.size = 0
self.pointer = 0
def add_transition(self, transition):
"""Add a transition to the replay buffer."""
def set_idx(buffer, new_element):
buffer[self.pointer] = new_element
jax.tree_util.tree_map(set_idx, self._dict, transition)
self.pointer = (self.pointer + 1) % self.max_size
self.size = max(self.pointer, self.size)
def clear(self):
"""Clear the replay buffer."""
self.size = self.pointer = 0
@dataclasses.dataclass
class GCDataset:
"""Dataset class for goal-conditioned RL.
This class provides a method to sample a batch of transitions with goals (value_goals and actor_goals) from the
dataset. The goals are sampled from the current state, future states in the same trajectory, and random states.
It also supports frame stacking and random-cropping image augmentation.
It reads the following keys from the config:
- discount: Discount factor for geometric sampling.
- value_p_curgoal: Probability of using the current state as the value goal.
- value_p_trajgoal: Probability of using a future state in the same trajectory as the value goal.
- value_p_randomgoal: Probability of using a random state as the value goal.
- value_geom_sample: Whether to use geometric sampling for future value goals.
- actor_p_curgoal: Probability of using the current state as the actor goal.
- actor_p_trajgoal: Probability of using a future state in the same trajectory as the actor goal.
- actor_p_randomgoal: Probability of using a random state as the actor goal.
- actor_geom_sample: Whether to use geometric sampling for future actor goals.
- gc_negative: Whether to use '0 if s == g else -1' (True) or '1 if s == g else 0' (False) as the reward.
- p_aug: Probability of applying image augmentation.
- frame_stack: Number of frames to stack.
Attributes:
dataset: Dataset object.
config: Configuration dictionary.
preprocess_frame_stack: Whether to preprocess frame stacks. If False, frame stacks are computed on-the-fly. This
saves memory but may slow down training.
"""
dataset: Dataset
config: Any
preprocess_frame_stack: bool = True
def __post_init__(self):
self.size = self.dataset.size
# Pre-compute trajectory boundaries.
(self.terminal_locs,) = np.nonzero(self.dataset['terminals'] > 0)
self.initial_locs = np.concatenate([[0], self.terminal_locs[:-1] + 1])
assert self.terminal_locs[-1] == self.size - 1
# Assert probabilities sum to 1.
assert np.isclose(
self.config['value_p_curgoal'] + self.config['value_p_trajgoal'] + self.config['value_p_randomgoal'], 1.0
)
assert np.isclose(
self.config['actor_p_curgoal'] + self.config['actor_p_trajgoal'] + self.config['actor_p_randomgoal'], 1.0
)
if self.config['frame_stack'] is not None:
# Only support compact (observation-only) datasets.
assert 'next_observations' not in self.dataset
if self.preprocess_frame_stack:
stacked_observations = self.get_stacked_observations(np.arange(self.size))
self.dataset = Dataset(self.dataset.copy(dict(observations=stacked_observations)))
def sample(self, batch_size, idxs=None, evaluation=False):
"""Sample a batch of transitions with goals.
This method samples a batch of transitions with goals (value_goals and actor_goals) from the dataset. They are
stored in the keys 'value_goals' and 'actor_goals', respectively. It also computes the 'rewards' and 'masks'
based on the indices of the goals.
Args:
batch_size: Batch size.
idxs: Indices of the transitions to sample. If None, random indices are sampled.
evaluation: Whether to sample for evaluation. If True, image augmentation is not applied.
"""
if idxs is None:
idxs = self.dataset.get_random_idxs(batch_size)
batch = self.dataset.sample(batch_size, idxs)
if self.config['frame_stack'] is not None:
batch['observations'] = self.get_observations(idxs)
batch['next_observations'] = self.get_observations(idxs + 1)
value_goal_idxs = self.sample_goals(
idxs,
self.config['value_p_curgoal'],
self.config['value_p_trajgoal'],
self.config['value_p_randomgoal'],
self.config['value_geom_sample'],
)
actor_goal_idxs = self.sample_goals(
idxs,
self.config['actor_p_curgoal'],
self.config['actor_p_trajgoal'],
self.config['actor_p_randomgoal'],
self.config['actor_geom_sample'],
)
batch['value_goals'] = self.get_observations(value_goal_idxs)
batch['actor_goals'] = self.get_observations(actor_goal_idxs)
successes = (idxs == value_goal_idxs).astype(float)
batch['masks'] = 1.0 - successes
batch['rewards'] = successes - (1.0 if self.config['gc_negative'] else 0.0)
if self.config['p_aug'] is not None and not evaluation:
if np.random.rand() < self.config['p_aug']:
self.augment(batch, ['observations', 'next_observations', 'value_goals', 'actor_goals'])
return batch
def sample_goals(self, idxs, p_curgoal, p_trajgoal, p_randomgoal, geom_sample):
"""Sample goals for the given indices."""
batch_size = len(idxs)
# Random goals.
random_goal_idxs = self.dataset.get_random_idxs(batch_size)
# Goals from the same trajectory (excluding the current state, unless it is the final state).
final_state_idxs = self.terminal_locs[np.searchsorted(self.terminal_locs, idxs)]
if geom_sample:
# Geometric sampling.
offsets = np.random.geometric(p=1 - self.config['discount'], size=batch_size) # in [1, inf)
traj_goal_idxs = np.minimum(idxs + offsets, final_state_idxs)
else:
# Uniform sampling.
distances = np.random.rand(batch_size) # in [0, 1)
traj_goal_idxs = np.round(
(np.minimum(idxs + 1, final_state_idxs) * distances + final_state_idxs * (1 - distances))
).astype(int)
if p_curgoal == 1.0:
goal_idxs = idxs
else:
goal_idxs = np.where(
np.random.rand(batch_size) < p_trajgoal / (1.0 - p_curgoal), traj_goal_idxs, random_goal_idxs
)
# Goals at the current state.
goal_idxs = np.where(np.random.rand(batch_size) < p_curgoal, idxs, goal_idxs)
return goal_idxs
def augment(self, batch, keys):
"""Apply image augmentation to the given keys."""
padding = 3
batch_size = len(batch[keys[0]])
crop_froms = np.random.randint(0, 2 * padding + 1, (batch_size, 2))
crop_froms = np.concatenate([crop_froms, np.zeros((batch_size, 1), dtype=np.int64)], axis=1)
for key in keys:
batch[key] = jax.tree_util.tree_map(
lambda arr: np.array(batched_random_crop(arr, crop_froms, padding)) if len(arr.shape) == 4 else arr,
batch[key],
)
def get_observations(self, idxs):
"""Return the observations for the given indices."""
if self.config['frame_stack'] is None or self.preprocess_frame_stack:
return jax.tree_util.tree_map(lambda arr: arr[idxs], self.dataset['observations'])
else:
return self.get_stacked_observations(idxs)
def get_stacked_observations(self, idxs):
"""Return the frame-stacked observations for the given indices."""
initial_state_idxs = self.initial_locs[np.searchsorted(self.initial_locs, idxs, side='right') - 1]
rets = []
for i in reversed(range(self.config['frame_stack'])):
cur_idxs = np.maximum(idxs - i, initial_state_idxs)
rets.append(jax.tree_util.tree_map(lambda arr: arr[cur_idxs], self.dataset['observations']))
return jax.tree_util.tree_map(lambda *args: np.concatenate(args, axis=-1), *rets)
@dataclasses.dataclass
class HGCDataset(GCDataset):
"""Dataset class for hierarchical goal-conditioned RL.
This class extends GCDataset to support high-level actor goals and prediction targets. It reads the following
additional key from the config:
- subgoal_steps: Subgoal steps (i.e., the number of steps to reach the low-level goal).
"""
def sample(self, batch_size, idxs=None, evaluation=False):
"""Sample a batch of transitions with goals.
This method samples a batch of transitions with goals from the dataset. The goals are stored in the keys
'value_goals', 'low_actor_goals', 'high_actor_goals', and 'high_actor_targets'. It also computes the 'rewards'
and 'masks' based on the indices of the goals.
Args:
batch_size: Batch size.
idxs: Indices of the transitions to sample. If None, random indices are sampled.
evaluation: Whether to sample for evaluation. If True, image augmentation is not applied.
"""
if idxs is None:
idxs = self.dataset.get_random_idxs(batch_size)
batch = self.dataset.sample(batch_size, idxs)
if self.config['frame_stack'] is not None:
batch['observations'] = self.get_observations(idxs)
batch['next_observations'] = self.get_observations(idxs + 1)
# Sample value goals.
value_goal_idxs = self.sample_goals(
idxs,
self.config['value_p_curgoal'],
self.config['value_p_trajgoal'],
self.config['value_p_randomgoal'],
self.config['value_geom_sample'],
)
batch['value_goals'] = self.get_observations(value_goal_idxs)
successes = (idxs == value_goal_idxs).astype(float)
batch['masks'] = 1.0 - successes
batch['rewards'] = successes - (1.0 if self.config['gc_negative'] else 0.0)
# Set low-level actor goals.
final_state_idxs = self.terminal_locs[np.searchsorted(self.terminal_locs, idxs)]
low_goal_idxs = np.minimum(idxs + self.config['subgoal_steps'], final_state_idxs)
batch['low_actor_goals'] = self.get_observations(low_goal_idxs)
# Sample high-level actor goals and set prediction targets.
# High-level future goals.
if self.config['actor_geom_sample']:
# Geometric sampling.
offsets = np.random.geometric(p=1 - self.config['discount'], size=batch_size) # in [1, inf)
high_traj_goal_idxs = np.minimum(idxs + offsets, final_state_idxs)
else:
# Uniform sampling.
distances = np.random.rand(batch_size) # in [0, 1)
high_traj_goal_idxs = np.round(
(np.minimum(idxs + 1, final_state_idxs) * distances + final_state_idxs * (1 - distances))
).astype(int)
high_traj_target_idxs = np.minimum(idxs + self.config['subgoal_steps'], high_traj_goal_idxs)
# High-level random goals.
high_random_goal_idxs = self.dataset.get_random_idxs(batch_size)
high_random_target_idxs = np.minimum(idxs + self.config['subgoal_steps'], final_state_idxs)
# Pick between high-level future goals and random goals.
pick_random = np.random.rand(batch_size) < self.config['actor_p_randomgoal']
high_goal_idxs = np.where(pick_random, high_random_goal_idxs, high_traj_goal_idxs)
high_target_idxs = np.where(pick_random, high_random_target_idxs, high_traj_target_idxs)
batch['high_actor_goals'] = self.get_observations(high_goal_idxs)
batch['high_actor_targets'] = self.get_observations(high_target_idxs)
if self.config['p_aug'] is not None and not evaluation:
if np.random.rand() < self.config['p_aug']:
self.augment(
batch,
[
'observations',
'next_observations',
'value_goals',
'low_actor_goals',
'high_actor_goals',
'high_actor_targets',
],
)
return batch
================================================
FILE: impls/utils/encoders.py
================================================
import functools
from typing import Sequence
import flax.linen as nn
import jax.numpy as jnp
from utils.networks import MLP
class ResnetStack(nn.Module):
"""ResNet stack module."""
num_features: int
num_blocks: int
max_pooling: bool = True
@nn.compact
def __call__(self, x):
initializer = nn.initializers.xavier_uniform()
conv_out = nn.Conv(
features=self.num_features,
kernel_size=(3, 3),
strides=1,
kernel_init=initializer,
padding='SAME',
)(x)
if self.max_pooling:
conv_out = nn.max_pool(
conv_out,
window_shape=(3, 3),
padding='SAME',
strides=(2, 2),
)
for _ in range(self.num_blocks):
block_input = conv_out
conv_out = nn.relu(conv_out)
conv_out = nn.Conv(
features=self.num_features,
kernel_size=(3, 3),
strides=1,
padding='SAME',
kernel_init=initializer,
)(conv_out)
conv_out = nn.relu(conv_out)
conv_out = nn.Conv(
features=self.num_features,
kernel_size=(3, 3),
strides=1,
padding='SAME',
kernel_init=initializer,
)(conv_out)
conv_out += block_input
return conv_out
class ImpalaEncoder(nn.Module):
"""IMPALA encoder."""
width: int = 1
stack_sizes: tuple = (16, 32, 32)
num_blocks: int = 2
dropout_rate: float = None
mlp_hidden_dims: Sequence[int] = (512,)
layer_norm: bool = False
def setup(self):
stack_sizes = self.stack_sizes
self.stack_blocks = [
ResnetStack(
num_features=stack_sizes[i] * self.width,
num_blocks=self.num_blocks,
)
for i in range(len(stack_sizes))
]
if self.dropout_rate is not None:
self.dropout = nn.Dropout(rate=self.dropout_rate)
@nn.compact
def __call__(self, x, train=True, cond_var=None):
x = x.astype(jnp.float32) / 255.0
conv_out = x
for idx in range(len(self.stack_blocks)):
conv_out = self.stack_blocks[idx](conv_out)
if self.dropout_rate is not None:
conv_out = self.dropout(conv_out, deterministic=not train)
conv_out = nn.relu(conv_out)
if self.layer_norm:
conv_out = nn.LayerNorm()(conv_out)
out = conv_out.reshape((*x.shape[:-3], -1))
out = MLP(self.mlp_hidden_dims, activate_final=True, layer_norm=self.layer_norm)(out)
return out
class GCEncoder(nn.Module):
"""Helper module to handle inputs to goal-conditioned networks.
It takes in observations (s) and goals (g) and returns the concatenation of `state_encoder(s)`, `goal_encoder(g)`,
and `concat_encoder([s, g])`. It ignores the encoders that are not provided. This way, the module can handle both
early and late fusion (or their variants) of state and goal information.
"""
state_encoder: nn.Module = None
goal_encoder: nn.Module = None
concat_encoder: nn.Module = None
@nn.compact
def __call__(self, observations, goals=None, goal_encoded=False):
"""Returns the representations of observations and goals.
If `goal_encoded` is True, `goals` is assumed to be already encoded representations. In this case, either
`goal_encoder` or `concat_encoder` must be None.
"""
reps = []
if self.state_encoder is not None:
reps.append(self.state_encoder(observations))
if goals is not None:
if goal_encoded:
# Can't have both goal_encoder and concat_encoder in this case.
assert self.goal_encoder is None or self.concat_encoder is None
reps.append(goals)
else:
if self.goal_encoder is not None:
reps.append(self.goal_encoder(goals))
if self.concat_encoder is not None:
reps.append(self.concat_encoder(jnp.concatenate([observations, goals], axis=-1)))
reps = jnp.concatenate(reps, axis=-1)
return reps
encoder_modules = {
'impala': ImpalaEncoder,
'impala_debug': functools.partial(ImpalaEncoder, num_blocks=1, stack_sizes=(4, 4)),
'impala_small': functools.partial(ImpalaEncoder, num_blocks=1),
'impala_large': functools.partial(ImpalaEncoder, stack_sizes=(64, 128, 128), mlp_hidden_dims=(1024,)),
}
================================================
FILE: impls/utils/env_utils.py
================================================
import collections
import os
import platform
import time
import gymnasium
import numpy as np
from gymnasium.spaces import Box
import ogbench
from utils.datasets import Dataset
class EpisodeMonitor(gymnasium.Wrapper):
"""Environment wrapper to monitor episode statistics."""
def __init__(self, env):
super().__init__(env)
self._reset_stats()
self.total_timesteps = 0
def _reset_stats(self):
self.reward_sum = 0.0
self.episode_length = 0
self.start_time = time.time()
def step(self, action):
observation, reward, terminated, truncated, info = self.env.step(action)
self.reward_sum += reward
self.episode_length += 1
self.total_timesteps += 1
info['total'] = {'timesteps': self.total_timesteps}
if terminated or truncated:
info['episode'] = {}
info['episode']['return'] = self.reward_sum
info['episode']['length'] = self.episode_length
info['episode']['duration'] = time.time() - self.start_time
return observation, reward, terminated, truncated, info
def reset(self, *args, **kwargs):
self._reset_stats()
return self.env.reset(*args, **kwargs)
class FrameStackWrapper(gymnasium.Wrapper):
"""Environment wrapper to stack observations."""
def __init__(self, env, num_stack):
super().__init__(env)
self.num_stack = num_stack
self.frames = collections.deque(maxlen=num_stack)
low = np.concatenate([self.observation_space.low] * num_stack, axis=-1)
high = np.concatenate([self.observation_space.high] * num_stack, axis=-1)
self.observation_space = Box(low=low, high=high, dtype=self.observation_space.dtype)
def get_observation(self):
assert len(self.frames) == self.num_stack
return np.concatenate(list(self.frames), axis=-1)
def reset(self, **kwargs):
ob, info = self.env.reset(**kwargs)
for _ in range(self.num_stack):
self.frames.append(ob)
if 'goal' in info:
info['goal'] = np.concatenate([info['goal']] * self.num_stack, axis=-1)
return self.get_observation(), info
def step(self, action):
ob, reward, terminated, truncated, info = self.env.step(action)
self.frames.append(ob)
return self.get_observation(), reward, terminated, truncated, info
def make_env_and_datasets(dataset_name, frame_stack=None):
"""Make OGBench environment and datasets.
Args:
dataset_name: Name of the dataset.
frame_stack: Number of frames to stack.
Returns:
A tuple of the environment, training dataset, and validation dataset.
"""
# Use compact dataset to save memory.
env, train_dataset, val_dataset = ogbench.make_env_and_datasets(dataset_name, compact_dataset=True)
train_dataset = Dataset.create(**train_dataset)
val_dataset = Dataset.create(**val_dataset)
if frame_stack is not None:
env = FrameStackWrapper(env, frame_stack)
env.reset()
return env, train_dataset, val_dataset
================================================
FILE: impls/utils/evaluation.py
================================================
from collections import defaultdict
import jax
import numpy as np
from tqdm import trange
def supply_rng(f, rng=jax.random.PRNGKey(0)):
"""Helper function to split the random number generator key before each call to the function."""
def wrapped(*args, **kwargs):
nonlocal rng
rng, key = jax.random.split(rng)
return f(*args, seed=key, **kwargs)
return wrapped
def flatten(d, parent_key='', sep='.'):
"""Flatten a dictionary."""
items = []
for k, v in d.items():
new_key = parent_key + sep + k if parent_key else k
if hasattr(v, 'items'):
items.extend(flatten(v, new_key, sep=sep).items())
else:
items.append((new_key, v))
return dict(items)
def add_to(dict_of_lists, single_dict):
"""Append values to the corresponding lists in the dictionary."""
for k, v in single_dict.items():
dict_of_lists[k].append(v)
def evaluate(
agent,
env,
task_id=None,
config=None,
num_eval_episodes=50,
num_video_episodes=0,
video_frame_skip=3,
eval_temperature=0,
eval_gaussian=None,
):
"""Evaluate the agent in the environment.
Args:
agent: Agent.
env: Environment.
task_id: Task ID to be passed to the environment.
config: Configuration dictionary.
num_eval_episodes: Number of episodes to evaluate the agent.
num_video_episodes: Number of episodes to render. These episodes are not included in the statistics.
video_frame_skip: Number of frames to skip between renders.
eval_temperature: Action sampling temperature.
eval_gaussian: Standard deviation of the Gaussian noise to add to the actions.
Returns:
A tuple containing the statistics, trajectories, and rendered videos.
"""
actor_fn = supply_rng(agent.sample_actions, rng=jax.random.PRNGKey(np.random.randint(0, 2**32)))
trajs = []
stats = defaultdict(list)
renders = []
for i in trange(num_eval_episodes + num_video_episodes):
traj = defaultdict(list)
should_render = i >= num_eval_episodes
observation, info = env.reset(options=dict(task_id=task_id, render_goal=should_render))
goal = info.get('goal')
goal_frame = info.get('goal_rendered')
done = False
step = 0
render = []
while not done:
action = actor_fn(observations=observation, goals=goal, temperature=eval_temperature)
action = np.array(action)
if not config.get('discrete'):
if eval_gaussian is not None:
action = np.random.normal(action, eval_gaussian)
action = np.clip(action, -1, 1)
next_observation, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
step += 1
if should_render and (step % video_frame_skip == 0 or done):
frame = env.render().copy()
if goal_frame is not None:
render.append(np.concatenate([goal_frame, frame], axis=0))
else:
render.append(frame)
transition = dict(
observation=observation,
next_observation=next_observation,
action=action,
reward=reward,
done=done,
info=info,
)
add_to(traj, transition)
observation = next_observation
if i < num_eval_episodes:
add_to(stats, flatten(info))
trajs.append(traj)
else:
renders.append(np.array(render))
for k, v in stats.items():
stats[k] = np.mean(v)
return stats, trajs, renders
================================================
FILE: impls/utils/flax_utils.py
================================================
import functools
import glob
import os
import pickle
from typing import Any, Dict, Mapping, Sequence
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax
nonpytree_field = functools.partial(flax.struct.field, pytree_node=False)
class ModuleDict(nn.Module):
"""A dictionary of modules.
This allows sharing parameters between modules and provides a convenient way to access them.
Attributes:
modules: Dictionary of modules.
"""
modules: Dict[str, nn.Module]
@nn.compact
def __call__(self, *args, name=None, **kwargs):
"""Forward pass.
For initialization, call with `name=None` and provide the arguments for each module in `kwargs`.
Otherwise, call with `name=` and provide the arguments for that module.
"""
if name is None:
if kwargs.keys() != self.modules.keys():
raise ValueError(
f'When `name` is not specified, kwargs must contain the arguments for each module. '
f'Got kwargs keys {kwargs.keys()} but module keys {self.modules.keys()}'
)
out = {}
for key, value in kwargs.items():
if isinstance(value, Mapping):
out[key] = self.modules[key](**value)
elif isinstance(value, Sequence):
out[key] = self.modules[key](*value)
else:
out[key] = self.modules[key](value)
return out
return self.modules[name](*args, **kwargs)
class TrainState(flax.struct.PyTreeNode):
"""Custom train state for models.
Attributes:
step: Counter to keep track of the training steps. It is incremented by 1 after each `apply_gradients` call.
apply_fn: Apply function of the model.
model_def: Model definition.
params: Parameters of the model.
tx: optax optimizer.
opt_state: Optimizer state.
"""
step: int
apply_fn: Any = nonpytree_field()
model_def: Any = nonpytree_field()
params: Any
tx: Any = nonpytree_field()
opt_state: Any
@classmethod
def create(cls, model_def, params, tx=None, **kwargs):
"""Create a new train state."""
if tx is not None:
opt_state = tx.init(params)
else:
opt_state = None
return cls(
step=1,
apply_fn=model_def.apply,
model_def=model_def,
params=params,
tx=tx,
opt_state=opt_state,
**kwargs,
)
def __call__(self, *args, params=None, method=None, **kwargs):
"""Forward pass.
When `params` is not provided, it uses the stored parameters.
The typical use case is to set `params` to `None` when you want to *stop* the gradients, and to pass the current
traced parameters when you want to flow the gradients. In other words, the default behavior is to stop the
gradients, and you need to explicitly provide the parameters to flow the gradients.
Args:
*args: Arguments to pass to the model.
params: Parameters to use for the forward pass. If `None`, it uses the stored parameters, without flowing
the gradients.
method: Method to call in the model. If `None`, it uses the default `apply` method.
**kwargs: Keyword arguments to pass to the model.
"""
if params is None:
params = self.params
variables = {'params': params}
if method is not None:
method_name = getattr(self.model_def, method)
else:
method_name = None
return self.apply_fn(variables, *args, method=method_name, **kwargs)
def select(self, name):
"""Helper function to select a module from a `ModuleDict`."""
return functools.partial(self, name=name)
def apply_gradients(self, grads, **kwargs):
"""Apply the gradients and return the updated state."""
updates, new_opt_state = self.tx.update(grads, self.opt_state, self.params)
new_params = optax.apply_updates(self.params, updates)
return self.replace(
step=self.step + 1,
params=new_params,
opt_state=new_opt_state,
**kwargs,
)
def apply_loss_fn(self, loss_fn):
"""Apply the loss function and return the updated state and info.
It additionally computes the gradient statistics and adds them to the dictionary.
"""
grads, info = jax.grad(loss_fn, has_aux=True)(self.params)
grad_max = jax.tree_util.tree_map(jnp.max, grads)
grad_min = jax.tree_util.tree_map(jnp.min, grads)
grad_norm = jax.tree_util.tree_map(jnp.linalg.norm, grads)
grad_max_flat = jnp.concatenate([jnp.reshape(x, -1) for x in jax.tree_util.tree_leaves(grad_max)], axis=0)
grad_min_flat = jnp.concatenate([jnp.reshape(x, -1) for x in jax.tree_util.tree_leaves(grad_min)], axis=0)
grad_norm_flat = jnp.concatenate([jnp.reshape(x, -1) for x in jax.tree_util.tree_leaves(grad_norm)], axis=0)
final_grad_max = jnp.max(grad_max_flat)
final_grad_min = jnp.min(grad_min_flat)
final_grad_norm = jnp.linalg.norm(grad_norm_flat, ord=1)
info.update(
{
'grad/max': final_grad_max,
'grad/min': final_grad_min,
'grad/norm': final_grad_norm,
}
)
return self.apply_gradients(grads=grads), info
def save_agent(agent, save_dir, epoch):
"""Save the agent to a file.
Args:
agent: Agent.
save_dir: Directory to save the agent.
epoch: Epoch number.
"""
save_dict = dict(
agent=flax.serialization.to_state_dict(agent),
)
save_path = os.path.join(save_dir, f'params_{epoch}.pkl')
with open(save_path, 'wb') as f:
pickle.dump(save_dict, f)
print(f'Saved to {save_path}')
def restore_agent(agent, restore_path, restore_epoch):
"""Restore the agent from a file.
Args:
agent: Agent.
restore_path: Path to the directory containing the saved agent.
restore_epoch: Epoch number.
"""
candidates = glob.glob(restore_path)
assert len(candidates) == 1, f'Found {len(candidates)} candidates: {candidates}'
restore_path = candidates[0] + f'/params_{restore_epoch}.pkl'
with open(restore_path, 'rb') as f:
load_dict = pickle.load(f)
agent = flax.serialization.from_state_dict(agent, load_dict['agent'])
print(f'Restored from {restore_path}')
return agent
================================================
FILE: impls/utils/log_utils.py
================================================
import os
import tempfile
from datetime import datetime
import absl.flags as flags
import ml_collections
import numpy as np
import wandb
from PIL import Image, ImageEnhance
class CsvLogger:
"""CSV logger for logging metrics to a CSV file."""
def __init__(self, path):
self.path = path
self.header = None
self.file = None
self.disallowed_types = (wandb.Image, wandb.Video, wandb.Histogram)
def log(self, row, step):
row['step'] = step
if self.file is None:
self.file = open(self.path, 'w')
if self.header is None:
self.header = [k for k, v in row.items() if not isinstance(v, self.disallowed_types)]
self.file.write(','.join(self.header) + '\n')
filtered_row = {k: v for k, v in row.items() if not isinstance(v, self.disallowed_types)}
self.file.write(','.join([str(filtered_row.get(k, '')) for k in self.header]) + '\n')
else:
filtered_row = {k: v for k, v in row.items() if not isinstance(v, self.disallowed_types)}
self.file.write(','.join([str(filtered_row.get(k, '')) for k in self.header]) + '\n')
self.file.flush()
def close(self):
if self.file is not None:
self.file.close()
def get_exp_name(seed):
"""Return the experiment name."""
exp_name = ''
exp_name += f'sd{seed:03d}_'
if 'SLURM_JOB_ID' in os.environ:
exp_name += f's_{os.environ["SLURM_JOB_ID"]}.'
if 'SLURM_PROCID' in os.environ:
exp_name += f'{os.environ["SLURM_PROCID"]}.'
exp_name += f'{datetime.now().strftime("%Y%m%d_%H%M%S")}'
return exp_name
def get_flag_dict():
"""Return the dictionary of flags."""
flag_dict = {k: getattr(flags.FLAGS, k) for k in flags.FLAGS if '.' not in k}
for k in flag_dict:
if isinstance(flag_dict[k], ml_collections.ConfigDict):
flag_dict[k] = flag_dict[k].to_dict()
return flag_dict
def setup_wandb(
entity=None,
project='project',
group=None,
name=None,
mode='online',
):
"""Set up Weights & Biases for logging."""
wandb_output_dir = tempfile.mkdtemp()
tags = [group] if group is not None else None
init_kwargs = dict(
config=get_flag_dict(),
project=project,
entity=entity,
tags=tags,
group=group,
dir=wandb_output_dir,
name=name,
settings=wandb.Settings(
start_method='thread',
_disable_stats=False,
),
mode=mode,
save_code=True,
)
run = wandb.init(**init_kwargs)
return run
def reshape_video(v, n_cols=None):
"""Helper function to reshape videos."""
if v.ndim == 4:
v = v[None,]
_, t, h, w, c = v.shape
if n_cols is None:
# Set n_cols to the square root of the number of videos.
n_cols = np.ceil(np.sqrt(v.shape[0])).astype(int)
if v.shape[0] % n_cols != 0:
len_addition = n_cols - v.shape[0] % n_cols
v = np.concatenate((v, np.zeros(shape=(len_addition, t, h, w, c))), axis=0)
n_rows = v.shape[0] // n_cols
v = np.reshape(v, (n_rows, n_cols, t, h, w, c))
v = np.transpose(v, axes=(2, 5, 0, 3, 1, 4))
v = np.reshape(v, (t, c, n_rows * h, n_cols * w))
return v
def get_wandb_video(renders=None, n_cols=None, fps=15):
"""Return a Weights & Biases video.
It takes a list of videos and reshapes them into a single video with the specified number of columns.
Args:
renders: List of videos. Each video should be a numpy array of shape (t, h, w, c).
n_cols: Number of columns for the reshaped video. If None, it is set to the square root of the number of videos.
"""
# Pad videos to the same length.
max_length = max([len(render) for render in renders])
for i, render in enumerate(renders):
assert render.dtype == np.uint8
# Decrease brightness of the padded frames.
final_frame = render[-1]
final_image = Image.fromarray(final_frame)
enhancer = ImageEnhance.Brightness(final_image)
final_image = enhancer.enhance(0.5)
final_frame = np.array(final_image)
pad = np.repeat(final_frame[np.newaxis, ...], max_length - len(render), axis=0)
renders[i] = np.concatenate([render, pad], axis=0)
# Add borders.
renders[i] = np.pad(renders[i], ((0, 0), (1, 1), (1, 1), (0, 0)), mode='constant', constant_values=0)
renders = np.array(renders) # (n, t, h, w, c)
renders = reshape_video(renders, n_cols) # (t, c, nr * h, nc * w)
return wandb.Video(renders, fps=fps, format='mp4')
================================================
FILE: impls/utils/networks.py
================================================
from typing import Any, Optional, Sequence
import distrax
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
def default_init(scale=1.0):
"""Default kernel initializer."""
return nn.initializers.variance_scaling(scale, 'fan_avg', 'uniform')
def ensemblize(cls, num_qs, out_axes=0, **kwargs):
"""Ensemblize a module."""
return nn.vmap(
cls,
variable_axes={'params': 0},
split_rngs={'params': True},
in_axes=None,
out_axes=out_axes,
axis_size=num_qs,
**kwargs,
)
class Identity(nn.Module):
"""Identity layer."""
def __call__(self, x):
return x
class MLP(nn.Module):
"""Multi-layer perceptron.
Attributes:
hidden_dims: Hidden layer dimensions.
activations: Activation function.
activate_final: Whether to apply activation to the final layer.
kernel_init: Kernel initializer.
layer_norm: Whether to apply layer normalization.
"""
hidden_dims: Sequence[int]
activations: Any = nn.gelu
activate_final: bool = False
kernel_init: Any = default_init()
layer_norm: bool = False
@nn.compact
def __call__(self, x):
for i, size in enumerate(self.hidden_dims):
x = nn.Dense(size, kernel_init=self.kernel_init)(x)
if i + 1 < len(self.hidden_dims) or self.activate_final:
x = self.activations(x)
if self.layer_norm:
x = nn.LayerNorm()(x)
return x
class LengthNormalize(nn.Module):
"""Length normalization layer.
It normalizes the input along the last dimension to have a length of sqrt(dim).
"""
@nn.compact
def __call__(self, x):
return x / jnp.linalg.norm(x, axis=-1, keepdims=True) * jnp.sqrt(x.shape[-1])
class Param(nn.Module):
"""Scalar parameter module."""
init_value: float = 0.0
@nn.compact
def __call__(self):
return self.param('value', init_fn=lambda key: jnp.full((), self.init_value))
class LogParam(nn.Module):
"""Scalar parameter module with log scale."""
init_value: float = 1.0
@nn.compact
def __call__(self):
log_value = self.param('log_value', init_fn=lambda key: jnp.full((), jnp.log(self.init_value)))
return jnp.exp(log_value)
class TransformedWithMode(distrax.Transformed):
"""Transformed distribution with mode calculation."""
def mode(self):
return self.bijector.forward(self.distribution.mode())
class RunningMeanStd(flax.struct.PyTreeNode):
"""Running mean and standard deviation.
Attributes:
eps: Epsilon value to avoid division by zero.
mean: Running mean.
var: Running variance.
clip_max: Clip value after normalization.
count: Number of samples.
"""
eps: Any = 1e-6
mean: Any = 1.0
var: Any = 1.0
clip_max: Any = 10.0
count: int = 0
def normalize(self, batch):
batch = (batch - self.mean) / jnp.sqrt(self.var + self.eps)
batch = jnp.clip(batch, -self.clip_max, self.clip_max)
return batch
def unnormalize(self, batch):
return batch * jnp.sqrt(self.var + self.eps) + self.mean
def update(self, batch):
batch_mean, batch_var = jnp.mean(batch, axis=0), jnp.var(batch, axis=0)
batch_count = len(batch)
delta = batch_mean - self.mean
total_count = self.count + batch_count
new_mean = self.mean + delta * batch_count / total_count
m_a = self.var * self.count
m_b = batch_var * batch_count
m_2 = m_a + m_b + delta**2 * self.count * batch_count / total_count
new_var = m_2 / total_count
return self.replace(mean=new_mean, var=new_var, count=total_count)
class GCActor(nn.Module):
"""Goal-conditioned actor.
Attributes:
hidden_dims: Hidden layer dimensions.
action_dim: Action dimension.
log_std_min: Minimum value of log standard deviation.
log_std_max: Maximum value of log standard deviation.
tanh_squash: Whether to squash the action with tanh.
state_dependent_std: Whether to use state-dependent standard deviation.
const_std: Whether to use constant standard deviation.
final_fc_init_scale: Initial scale of the final fully-connected layer.
gc_encoder: Optional GCEncoder module to encode the inputs.
"""
hidden_dims: Sequence[int]
action_dim: int
log_std_min: Optional[float] = -5
log_std_max: Optional[float] = 2
tanh_squash: bool = False
state_dependent_std: bool = False
const_std: bool = True
final_fc_init_scale: float = 1e-2
gc_encoder: nn.Module = None
def setup(self):
self.actor_net = MLP(self.hidden_dims, activate_final=True)
self.mean_net = nn.Dense(self.action_dim, kernel_init=default_init(self.final_fc_init_scale))
if self.state_dependent_std:
self.log_std_net = nn.Dense(self.action_dim, kernel_init=default_init(self.final_fc_init_scale))
else:
if not self.const_std:
self.log_stds = self.param('log_stds', nn.initializers.zeros, (self.action_dim,))
def __call__(
self,
observations,
goals=None,
goal_encoded=False,
temperature=1.0,
):
"""Return the action distribution.
Args:
observations: Observations.
goals: Goals (optional).
goal_encoded: Whether the goals are already encoded.
temperature: Scaling factor for the standard deviation.
"""
if self.gc_encoder is not None:
inputs = self.gc_encoder(observations, goals, goal_encoded=goal_encoded)
else:
inputs = [observations]
if goals is not None:
inputs.append(goals)
inputs = jnp.concatenate(inputs, axis=-1)
outputs = self.actor_net(inputs)
means = self.mean_net(outputs)
if self.state_dependent_std:
log_stds = self.log_std_net(outputs)
else:
if self.const_std:
log_stds = jnp.zeros_like(means)
else:
log_stds = self.log_stds
log_stds = jnp.clip(log_stds, self.log_std_min, self.log_std_max)
distribution = distrax.MultivariateNormalDiag(loc=means, scale_diag=jnp.exp(log_stds) * temperature)
if self.tanh_squash:
distribution = TransformedWithMode(distribution, distrax.Block(distrax.Tanh(), ndims=1))
return distribution
class GCDiscreteActor(nn.Module):
"""Goal-conditioned actor for discrete actions.
Attributes:
hidden_dims: Hidden layer dimensions.
action_dim: Action dimension.
final_fc_init_scale: Initial scale of the final fully-connected layer.
gc_encoder: Optional GCEncoder module to encode the inputs.
"""
hidden_dims: Sequence[int]
action_dim: int
final_fc_init_scale: float = 1e-2
gc_encoder: nn.Module = None
def setup(self):
self.actor_net = MLP(self.hidden_dims, activate_final=True)
self.logit_net = nn.Dense(self.action_dim, kernel_init=default_init(self.final_fc_init_scale))
def __call__(
self,
observations,
goals=None,
goal_encoded=False,
temperature=1.0,
):
"""Return the action distribution.
Args:
observations: Observations.
goals: Goals (optional).
goal_encoded: Whether the goals are already encoded.
temperature: Inverse scaling factor for the logits (set to 0 to get the argmax).
"""
if self.gc_encoder is not None:
inputs = self.gc_encoder(observations, goals, goal_encoded=goal_encoded)
else:
inputs = [observations]
if goals is not None:
inputs.append(goals)
inputs = jnp.concatenate(inputs, axis=-1)
outputs = self.actor_net(inputs)
logits = self.logit_net(outputs)
distribution = distrax.Categorical(logits=logits / jnp.maximum(1e-6, temperature))
return distribution
class GCValue(nn.Module):
"""Goal-conditioned value/critic function.
This module can be used for both value V(s, g) and critic Q(s, a, g) functions.
Attributes:
hidden_dims: Hidden layer dimensions.
layer_norm: Whether to apply layer normalization.
ensemble: Whether to ensemble the value function.
gc_encoder: Optional GCEncoder module to encode the inputs.
"""
hidden_dims: Sequence[int]
layer_norm: bool = True
ensemble: bool = True
gc_encoder: nn.Module = None
def setup(self):
mlp_module = MLP
if self.ensemble:
mlp_module = ensemblize(mlp_module, 2)
value_net = mlp_module((*self.hidden_dims, 1), activate_final=False, layer_norm=self.layer_norm)
self.value_net = value_net
def __call__(self, observations, goals=None, actions=None):
"""Return the value/critic function.
Args:
observations: Observations.
goals: Goals (optional).
actions: Actions (optional).
"""
if self.gc_encoder is not None:
inputs = [self.gc_encoder(observations, goals)]
else:
inputs = [observations]
if goals is not None:
inputs.append(goals)
if actions is not None:
inputs.append(actions)
inputs = jnp.concatenate(inputs, axis=-1)
v = self.value_net(inputs).squeeze(-1)
return v
class GCDiscreteCritic(GCValue):
"""Goal-conditioned critic for discrete actions."""
action_dim: int = None
def __call__(self, observations, goals=None, actions=None):
actions = jnp.eye(self.action_dim)[actions]
return super().__call__(observations, goals, actions)
class GCBilinearValue(nn.Module):
"""Goal-conditioned bilinear value/critic function.
This module computes the value function as V(s, g) = phi(s)^T psi(g) / sqrt(d) or the critic function as
Q(s, a, g) = phi(s, a)^T psi(g) / sqrt(d), where phi and psi output d-dimensional vectors.
Attributes:
hidden_dims: Hidden layer dimensions.
latent_dim: Latent dimension.
layer_norm: Whether to apply layer normalization.
ensemble: Whether to ensemble the value function.
value_exp: Whether to exponentiate the value. Useful for contrastive learning.
state_encoder: Optional state encoder.
goal_encoder: Optional goal encoder.
"""
hidden_dims: Sequence[int]
latent_dim: int
layer_norm: bool = True
ensemble: bool = True
value_exp: bool = False
state_encoder: nn.Module = None
goal_encoder: nn.Module = None
def setup(self):
mlp_module = MLP
if self.ensemble:
mlp_module = ensemblize(mlp_module, 2)
self.phi = mlp_module((*self.hidden_dims, self.latent_dim), activate_final=False, layer_norm=self.layer_norm)
self.psi = mlp_module((*self.hidden_dims, self.latent_dim), activate_final=False, layer_norm=self.layer_norm)
def __call__(self, observations, goals, actions=None, info=False):
"""Return the value/critic function.
Args:
observations: Observations.
goals: Goals.
actions: Actions (optional).
info: Whether to additionally return the representations phi and psi.
"""
if self.state_encoder is not None:
observations = self.state_encoder(observations)
if self.goal_encoder is not None:
goals = self.goal_encoder(goals)
if actions is None:
phi_inputs = observations
else:
phi_inputs = jnp.concatenate([observations, actions], axis=-1)
phi = self.phi(phi_inputs)
psi = self.psi(goals)
v = (phi * psi / jnp.sqrt(self.latent_dim)).sum(axis=-1)
if self.value_exp:
v = jnp.exp(v)
if info:
return v, phi, psi
else:
return v
class GCDiscreteBilinearCritic(GCBilinearValue):
"""Goal-conditioned bilinear critic for discrete actions."""
action_dim: int = None
def __call__(self, observations, goals=None, actions=None, info=False):
actions = jnp.eye(self.action_dim)[actions]
return super().__call__(observations, goals, actions, info)
class GCMRNValue(nn.Module):
"""Metric residual network (MRN) value function.
This module computes the value function as the sum of a symmetric Euclidean distance and an asymmetric
L^infinity-based quasimetric.
Attributes:
hidden_dims: Hidden layer dimensions.
latent_dim: Latent dimension.
layer_norm: Whether to apply layer normalization.
encoder: Optional state/goal encoder.
"""
hidden_dims: Sequence[int]
latent_dim: int
layer_norm: bool = True
encoder: nn.Module = None
def setup(self):
self.phi = MLP((*self.hidden_dims, self.latent_dim), activate_final=False, layer_norm=self.layer_norm)
def __call__(self, observations, goals, is_phi=False, info=False):
"""Return the MRN value function.
Args:
observations: Observations.
goals: Goals.
is_phi: Whether the inputs are already encoded by phi.
info: Whether to additionally return the representations phi_s and phi_g.
"""
if is_phi:
phi_s = observations
phi_g = goals
else:
if self.encoder is not None:
observations = self.encoder(observations)
goals = self.encoder(goals)
phi_s = self.phi(observations)
phi_g = self.phi(goals)
sym_s = phi_s[..., : self.latent_dim // 2]
sym_g = phi_g[..., : self.latent_dim // 2]
asym_s = phi_s[..., self.latent_dim // 2 :]
asym_g = phi_g[..., self.latent_dim // 2 :]
squared_dist = ((sym_s - sym_g) ** 2).sum(axis=-1)
quasi = jax.nn.relu((asym_s - asym_g).max(axis=-1))
v = jnp.sqrt(jnp.maximum(squared_dist, 1e-12)) + quasi
if info:
return v, phi_s, phi_g
else:
return v
class GCIQEValue(nn.Module):
"""Interval quasimetric embedding (IQE) value function.
This module computes the value function as an IQE-based quasimetric.
Attributes:
hidden_dims: Hidden layer dimensions.
latent_dim: Latent dimension.
dim_per_component: Dimension of each component in IQE (i.e., number of intervals in each group).
layer_norm: Whether to apply layer normalization.
encoder: Optional state/goal encoder.
"""
hidden_dims: Sequence[int]
latent_dim: int
dim_per_component: int
layer_norm: bool = True
encoder: nn.Module = None
def setup(self):
self.phi = MLP((*self.hidden_dims, self.latent_dim), activate_final=False, layer_norm=self.layer_norm)
self.alpha = Param()
def __call__(self, observations, goals, is_phi=False, info=False):
"""Return the IQE value function.
Args:
observations: Observations.
goals: Goals.
is_phi: Whether the inputs are already encoded by phi.
info: Whether to additionally return the representations phi_s and phi_g.
"""
alpha = jax.nn.sigmoid(self.alpha())
if is_phi:
phi_s = observations
phi_g = goals
else:
if self.encoder is not None:
observations = self.encoder(observations)
goals = self.encoder(goals)
phi_s = self.phi(observations)
phi_g = self.phi(goals)
x = jnp.reshape(phi_s, (*phi_s.shape[:-1], -1, self.dim_per_component))
y = jnp.reshape(phi_g, (*phi_g.shape[:-1], -1, self.dim_per_component))
valid = x < y
xy = jnp.concatenate(jnp.broadcast_arrays(x, y), axis=-1)
ixy = xy.argsort(axis=-1)
sxy = jnp.take_along_axis(xy, ixy, axis=-1)
neg_inc_copies = jnp.take_along_axis(valid, ixy % self.dim_per_component, axis=-1) * jnp.where(
ixy < self.dim_per_component, -1, 1
)
neg_inp_copies = jnp.cumsum(neg_inc_copies, axis=-1)
neg_f = -1.0 * (neg_inp_copies < 0)
neg_incf = jnp.concatenate([neg_f[..., :1], neg_f[..., 1:] - neg_f[..., :-1]], axis=-1)
components = (sxy * neg_incf).sum(axis=-1)
v = alpha * components.mean(axis=-1) + (1 - alpha) * components.max(axis=-1)
if info:
return v, phi_s, phi_g
else:
return v
================================================
FILE: ogbench/__init__.py
================================================
"""OGBench: Benchmarking Offline Goal-Conditioned RL"""
import ogbench.locomaze
import ogbench.manipspace
import ogbench.powderworld
from ogbench.utils import download_datasets, load_dataset, make_env_and_datasets
__all__ = (
'locomaze',
'manipspace',
'powderworld',
'download_datasets',
'load_dataset',
'make_env_and_datasets',
)
================================================
FILE: ogbench/locomaze/__init__.py
================================================
from gymnasium.envs.registration import register
visual_dict = dict(
ob_type='pixels',
render_mode='rgb_array',
width=64,
height=64,
camera_name='back',
)
singletask_dict = dict(
add_noise_to_goal=False,
success_timing='pre',
)
# Environments for offline goal-conditioned RL.
register(
id='pointmaze-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='medium',
),
)
register(
id='pointmaze-large-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='large',
),
)
register(
id='pointmaze-giant-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='giant',
),
)
register(
id='pointmaze-teleport-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='teleport',
),
)
register(
id='antmaze-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='medium',
),
)
register(
id='visual-antmaze-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='medium',
**visual_dict,
),
)
register(
id='antmaze-large-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='large',
),
)
register(
id='visual-antmaze-large-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='large',
**visual_dict,
),
)
register(
id='antmaze-giant-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='giant',
),
)
register(
id='visual-antmaze-giant-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='giant',
**visual_dict,
),
)
register(
id='antmaze-teleport-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='teleport',
),
)
register(
id='visual-antmaze-teleport-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='teleport',
**visual_dict,
),
)
register(
id='humanoidmaze-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='medium',
),
)
register(
id='visual-humanoidmaze-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='medium',
**visual_dict,
),
)
register(
id='humanoidmaze-large-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='large',
),
)
register(
id='visual-humanoidmaze-large-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='large',
**visual_dict,
),
)
register(
id='humanoidmaze-giant-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=4000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='giant',
),
)
register(
id='visual-humanoidmaze-giant-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=4000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='giant',
**visual_dict,
),
)
register(
id='humanoidmaze-teleport-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='teleport',
),
)
register(
id='visual-humanoidmaze-teleport-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='teleport',
**visual_dict,
),
)
register(
id='antsoccer-arena-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='ball',
maze_type='arena',
),
)
register(
id='antsoccer-medium-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='ball',
maze_type='medium',
),
)
# Environments for reward-based single-task offline RL.
for task_id in [None, 1, 2, 3, 4, 5]:
task_suffix = '' if task_id is None else f'-task{task_id}'
reward_task_id = 0 if task_id is None else task_id # 0 means the default task.
register(
id=f'pointmaze-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='medium',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'pointmaze-large-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='large',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'pointmaze-giant-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='giant',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'pointmaze-teleport-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='point',
maze_env_type='maze',
maze_type='teleport',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antmaze-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='medium',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-antmaze-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='medium',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antmaze-large-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='large',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-antmaze-large-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='large',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antmaze-giant-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='giant',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-antmaze-giant-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='giant',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antmaze-teleport-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='teleport',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-antmaze-teleport-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='maze',
maze_type='teleport',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'humanoidmaze-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='medium',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-humanoidmaze-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='medium',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'humanoidmaze-large-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='large',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-humanoidmaze-large-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='large',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'humanoidmaze-giant-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=4000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='giant',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-humanoidmaze-giant-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=4000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='giant',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'humanoidmaze-teleport-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='teleport',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'visual-humanoidmaze-teleport-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=2000,
kwargs=dict(
loco_env_type='humanoid',
maze_env_type='maze',
maze_type='teleport',
**visual_dict,
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antsoccer-arena-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='ball',
maze_type='arena',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
register(
id=f'antsoccer-medium-singletask{task_suffix}-v0',
entry_point='ogbench.locomaze.maze:make_maze_env',
max_episode_steps=1000,
kwargs=dict(
loco_env_type='ant',
maze_env_type='ball',
maze_type='medium',
reward_task_id=reward_task_id,
**singletask_dict,
),
)
================================================
FILE: ogbench/locomaze/ant.py
================================================
import os
import gymnasium
import numpy as np
from gymnasium import utils
from gymnasium.envs.mujoco import MujocoEnv
from gymnasium.spaces import Box
class AntEnv(MujocoEnv, utils.EzPickle):
"""Gymnasium Ant environment.
Unlike the original Ant environment, this environment uses a restricted joint range for the actuators, as typically
done in previous works in hierarchical reinforcement learning. It also uses a control frequency of 10Hz instead of
20Hz, which is the default in the original environment.
"""
xml_file = os.path.join(os.path.dirname(__file__), 'assets', 'ant.xml')
metadata = {
'render_modes': ['human', 'rgb_array', 'depth_array'],
'render_fps': 10,
}
if gymnasium.__version__ >= '1.1.0':
metadata['render_modes'] += ['rgbd_tuple']
def __init__(
self,
xml_file=None,
reset_noise_scale=0.1,
render_mode='rgb_array',
width=200,
height=200,
**kwargs,
):
"""Initialize the Ant environment.
Args:
xml_file: Path to the XML description (optional).
reset_noise_scale: Scale of the noise added to the initial state during reset.
render_mode: Rendering mode.
width: Width of the rendered image.
height: Height of the rendered image.
**kwargs: Additional keyword arguments.
"""
if xml_file is None:
xml_file = self.xml_file
utils.EzPickle.__init__(
self,
xml_file,
reset_noise_scale,
**kwargs,
)
self._reset_noise_scale = reset_noise_scale
observation_space = Box(low=-np.inf, high=np.inf, shape=(29,), dtype=np.float64)
MujocoEnv.__init__(
self,
xml_file,
frame_skip=5,
observation_space=observation_space,
render_mode=render_mode,
width=width,
height=height,
**kwargs,
)
def step(self, action):
prev_qpos = self.data.qpos.copy()
prev_qvel = self.data.qvel.copy()
self.do_simulation(action, self.frame_skip)
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
observation = self.get_ob()
if self.render_mode == 'human':
self.render()
return (
observation,
0.0,
False,
False,
{
'xy': self.get_xy(),
'prev_qpos': prev_qpos,
'prev_qvel': prev_qvel,
'qpos': qpos,
'qvel': qvel,
},
)
def get_ob(self):
position = self.data.qpos.flat.copy()
velocity = self.data.qvel.flat.copy()
return np.concatenate([position, velocity])
def reset_model(self):
noise_low = -self._reset_noise_scale
noise_high = self._reset_noise_scale
qpos = self.init_qpos + self.np_random.uniform(low=noise_low, high=noise_high, size=self.model.nq)
qvel = self.init_qvel + self._reset_noise_scale * self.np_random.standard_normal(self.model.nv)
self.set_state(qpos, qvel)
observation = self.get_ob()
return observation
def get_xy(self):
return self.data.qpos[:2].copy()
def set_xy(self, xy):
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
qpos[:2] = xy
self.set_state(qpos, qvel)
================================================
FILE: ogbench/locomaze/assets/ant.xml
================================================
================================================
FILE: ogbench/locomaze/assets/humanoid.xml
================================================
================================================
FILE: ogbench/locomaze/assets/point.xml
================================================
================================================
FILE: ogbench/locomaze/humanoid.py
================================================
import contextlib
import os
import gymnasium
import mujoco
import numpy as np
from gymnasium import utils
from gymnasium.envs.mujoco import MujocoEnv
from gymnasium.spaces import Box
class HumanoidEnv(MujocoEnv, utils.EzPickle):
"""DMC Humanoid environment.
Several methods are reimplemented to remove the dependency on the `dm_control` package. It is supposed to work
identically to the original Humanoid environment.
"""
xml_file = os.path.join(os.path.dirname(__file__), 'assets', 'humanoid.xml')
metadata = {
'render_modes': ['human', 'rgb_array', 'depth_array'],
'render_fps': 40,
}
if gymnasium.__version__ >= '1.1.0':
metadata['render_modes'] += ['rgbd_tuple']
def __init__(
self,
xml_file=None,
render_mode='rgb_array',
width=200,
height=200,
**kwargs,
):
"""Initialize the Humanoid environment.
Args:
xml_file: Path to the XML description (optional).
render_mode: Rendering mode.
width: Width of the rendered image.
height: Height of the rendered image.
**kwargs: Additional keyword arguments.
"""
if xml_file is None:
xml_file = self.xml_file
utils.EzPickle.__init__(
self,
xml_file,
**kwargs,
)
observation_space = Box(low=-np.inf, high=np.inf, shape=(69,), dtype=np.float64)
MujocoEnv.__init__(
self,
xml_file,
frame_skip=5,
observation_space=observation_space,
render_mode=render_mode,
width=width,
height=height,
**kwargs,
)
def step(self, action):
prev_qpos = self.data.qpos.copy()
prev_qvel = self.data.qvel.copy()
self.do_simulation(action, self.frame_skip)
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
observation = self.get_ob()
if self.render_mode == 'human':
self.render()
return (
observation,
0.0,
False,
False,
{
'xy': self.get_xy(),
'prev_qpos': prev_qpos,
'prev_qvel': prev_qvel,
'qpos': qpos,
'qvel': qvel,
},
)
def _step_mujoco_simulation(self, ctrl, n_frames):
self.data.ctrl[:] = ctrl
# DMC-style stepping (see Page 6 of https://arxiv.org/abs/2006.12983).
if self.model.opt.integrator != mujoco.mjtIntegrator.mjINT_RK4.value:
mujoco.mj_step2(self.model, self.data)
if n_frames > 1:
mujoco.mj_step(self.model, self.data, n_frames - 1)
else:
mujoco.mj_step(self.model, self.data, n_frames)
mujoco.mj_step1(self.model, self.data)
def get_ob(self):
xy = self.data.qpos[:2]
joint_angles = self.data.qpos[7:] # Skip the 7 DoFs of the free root joint.
head_height = self.data.xpos[2, 2] # ['head', 'z']
torso_frame = self.data.xmat[1].reshape(3, 3) # ['torso']
torso_pos = self.data.xpos[1] # ['torso']
positions = []
for idx in [16, 10, 13, 7]: # ['left_hand', 'left_foot', 'right_hand', 'right_foot']
torso_to_limb = self.data.xpos[idx] - torso_pos
positions.append(torso_to_limb.dot(torso_frame))
extremities = np.hstack(positions)
torso_vertical_orientation = self.data.xmat[1, [6, 7, 8]] # ['torso', ['zx', 'zy', 'zz']]
center_of_mass_velocity = self.data.sensordata[0:3] # ['torso_subtreelinvel']
velocity = self.data.qvel
return np.concatenate(
[
xy,
joint_angles,
[head_height],
extremities,
torso_vertical_orientation,
center_of_mass_velocity,
velocity,
]
)
@contextlib.contextmanager
def disable(self, *flags):
old_bitmask = self.model.opt.disableflags
new_bitmask = old_bitmask
for flag in flags:
if isinstance(flag, str):
field_name = 'mjDSBL_' + flag.upper()
flag = getattr(mujoco.mjtDisableBit, field_name)
elif isinstance(flag, int):
flag = mujoco.mjtDisableBit(flag)
new_bitmask |= flag.value
self.model.opt.disableflags = new_bitmask
try:
yield
finally:
self.model.opt.disableflags = old_bitmask
def reset_model(self):
penetrating = True
while penetrating:
quat = self.np_random.uniform(size=4)
quat /= np.linalg.norm(quat)
self.data.qpos[3:7] = quat
self.data.qvel = 0.1 * self.np_random.standard_normal(self.model.nv)
for joint_id in range(1, self.model.njnt):
range_min, range_max = self.model.jnt_range[joint_id]
self.data.qpos[6 + joint_id] = self.np_random.uniform(range_min, range_max)
with self.disable('actuation'):
mujoco.mj_forward(self.model, self.data)
penetrating = self.data.ncon > 0
observation = self.get_ob()
return observation
def get_xy(self):
return self.data.qpos[:2].copy()
def set_xy(self, xy):
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
qpos[:2] = xy
self.set_state(qpos, qvel)
================================================
FILE: ogbench/locomaze/maze.py
================================================
import tempfile
import xml.etree.ElementTree as ET
import mujoco
import numpy as np
from gymnasium.spaces import Box
from ogbench.locomaze.ant import AntEnv
from ogbench.locomaze.humanoid import HumanoidEnv
from ogbench.locomaze.point import PointEnv
def make_maze_env(loco_env_type, maze_env_type, *args, **kwargs):
"""Factory function for creating a maze environment.
Args:
loco_env_type: Locomotion environment type. One of 'point', 'ant', or 'humanoid'.
maze_env_type: Maze environment type. Either 'maze' or 'ball'.
*args: Additional arguments to pass to the target class.
**kwargs: Additional keyword arguments to pass to the target class.
"""
if loco_env_type == 'point':
loco_env_class = PointEnv
elif loco_env_type == 'ant':
loco_env_class = AntEnv
elif loco_env_type == 'humanoid':
loco_env_class = HumanoidEnv
else:
raise ValueError(f'Unknown locomotion environment type: {loco_env_type}')
class MazeEnv(loco_env_class):
"""Maze environment.
It inherits from the locomotion environment and adds a maze to it.
"""
def __init__(
self,
maze_type='large',
maze_unit=4.0,
maze_height=0.5,
terminate_at_goal=True,
success_timing='post',
ob_type='states',
add_noise_to_goal=True,
reward_task_id=None,
use_oracle_rep=False,
*args,
**kwargs,
):
"""Initialize the maze environment.
Args:
maze_type: Maze type. One of 'arena', 'medium', 'large', 'giant', or 'teleport'.
maze_unit: Size of a maze unit block.
maze_height: Height of the maze walls.
terminate_at_goal: Whether to terminate the episode when the goal is reached.
success_timing: When to compute the success and terminated flags and the reward. Either 'pre' (before
taking the action) or 'post' (after taking the action).
ob_type: Observation type. Either 'states' or 'pixels'.
add_noise_to_goal: Whether to add noise to the goal position.
reward_task_id: Task ID for single-task RL. If this is not None, the environment operates in a
single-task mode with the specified task ID. The task ID must be either a valid task ID or 0, where
0 means using the default task.
use_oracle_rep: Whether to use oracle goal representations.
*args: Additional arguments to pass to the parent locomotion environment.
**kwargs: Additional keyword arguments to pass to the parent locomotion environment.
"""
self._maze_type = maze_type
self._maze_unit = maze_unit
self._maze_height = maze_height
self._terminate_at_goal = terminate_at_goal
self._success_timing = success_timing
self._ob_type = ob_type
self._add_noise_to_goal = add_noise_to_goal
self._reward_task_id = reward_task_id
self._use_oracle_rep = use_oracle_rep
assert ob_type in ['states', 'pixels']
assert success_timing in ['pre', 'post']
# Define constants.
self._offset_x = 4
self._offset_y = 4
self._noise = 1
self._goal_tol = 1.0 if loco_env_type == 'point' else 0.5
# Define maze map.
self._teleport_info = None
if self._maze_type == 'arena':
maze_map = [
[1, 1, 1, 1, 1, 1, 1, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1],
[1, 1, 1, 1, 1, 1, 1, 1],
]
elif self._maze_type == 'medium':
maze_map = [
[1, 1, 1, 1, 1, 1, 1, 1],
[1, 0, 0, 1, 1, 0, 0, 1],
[1, 0, 0, 1, 0, 0, 0, 1],
[1, 1, 0, 0, 0, 1, 1, 1],
[1, 0, 0, 1, 0, 0, 0, 1],
[1, 0, 1, 0, 0, 1, 0, 1],
[1, 0, 0, 0, 1, 0, 0, 1],
[1, 1, 1, 1, 1, 1, 1, 1],
]
elif self._maze_type == 'large':
maze_map = [
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1],
[1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1],
[1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1],
[1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1],
[1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1],
[1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1],
[1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
]
elif self._maze_type == 'giant':
maze_map = [
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1],
[1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 1],
[1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
[1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1],
[1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1],
[1, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1],
[1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
[1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1],
[1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1],
[1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
]
elif self._maze_type == 'teleport':
maze_map = [
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1],
[1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1],
[1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1],
[1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1],
[1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1],
[1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1],
[1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
]
self._teleport_info = dict(
teleport_in_ijs=[(4, 6), (5, 1)],
teleport_out_ijs=[(1, 7), (6, 1), (6, 10)],
teleport_radius=1,
)
self._teleport_info['teleport_in_xys'] = [
self.ij_to_xy(ij) for ij in self._teleport_info['teleport_in_ijs']
]
self._teleport_info['teleport_out_xys'] = [
self.ij_to_xy(ij) for ij in self._teleport_info['teleport_out_ijs']
]
else:
raise ValueError(f'Unknown maze type: {self._maze_type}')
self.maze_map = np.array(maze_map)
# Update XML file.
xml_file = self.xml_file
tree = ET.parse(xml_file)
self.update_tree(tree)
_, maze_xml_file = tempfile.mkstemp(text=True, suffix='.xml')
tree.write(maze_xml_file)
super().__init__(xml_file=maze_xml_file, *args, **kwargs)
# Make custom camera.
if self.camera_id is None and self.camera_name is None:
# Use a custom default view.
camera = mujoco.MjvCamera()
camera.lookat[0] = 2 * (self.maze_map.shape[1] - 3)
camera.lookat[1] = 2 * (self.maze_map.shape[0] - 3)
camera.distance = 5 * (self.maze_map.shape[1] - 2)
camera.elevation = -90
self.custom_camera = camera
else:
self.custom_camera = self.camera_id or self.camera_name
# Set task goals.
self.task_infos = []
self.cur_task_id = None
self.cur_task_info = None
self.set_tasks()
self.num_tasks = len(self.task_infos)
self.cur_goal_xy = np.zeros(2)
self.custom_renderer = None
if self._ob_type == 'pixels':
self.observation_space = Box(low=0, high=255, shape=(64, 64, 3), dtype=np.uint8)
# Manually color the floor to enable the agent to infer its position from the observation.
tex_grid = self.model.tex('grid')
tex_height = tex_grid.height[0]
tex_width = tex_grid.width[0]
# MuJoCo 3.2.1 changed the attribute name from 'tex_rgb' to 'tex_data'.
attr_name = 'tex_rgb' if hasattr(self.model, 'tex_rgb') else 'tex_data'
tex_rgb = getattr(self.model, attr_name)[tex_grid.adr[0] : tex_grid.adr[0] + 3 * tex_height * tex_width]
tex_rgb = tex_rgb.reshape(tex_height, tex_width, 3)
for x in range(tex_height):
for y in range(tex_width):
min_value = 0
max_value = 192
r = int(x / tex_height * (max_value - min_value) + min_value)
g = int(y / tex_width * (max_value - min_value) + min_value)
tex_rgb[x, y, :] = [r, g, 128]
self.initialize_renderer()
else:
ex_ob = self.get_ob()
self.observation_space = Box(low=-np.inf, high=np.inf, shape=ex_ob.shape, dtype=ex_ob.dtype)
def update_tree(self, tree):
"""Update the XML tree to include the maze."""
worldbody = tree.find('.//worldbody')
# Add walls.
for i in range(self.maze_map.shape[0]):
for j in range(self.maze_map.shape[1]):
struct = self.maze_map[i, j]
if struct == 1:
ET.SubElement(
worldbody,
'geom',
name=f'block_{i}_{j}',
pos=f'{j * self._maze_unit - self._offset_x} {i * self._maze_unit - self._offset_y} {self._maze_height / 2 * self._maze_unit}',
size=f'{self._maze_unit / 2} {self._maze_unit / 2} {self._maze_height / 2 * self._maze_unit}',
type='box',
contype='1',
conaffinity='1',
material='wall',
)
# Adjust floor size.
center_x, center_y = 2 * (self.maze_map.shape[1] - 3), 2 * (self.maze_map.shape[0] - 3)
size_x, size_y = 2 * self.maze_map.shape[1], 2 * self.maze_map.shape[0]
floor = tree.find('.//geom[@name="floor"]')
floor.set('pos', f'{center_x} {center_y} 0')
floor.set('size', f'{size_x} {size_y} 0.2')
if self._teleport_info is not None:
# Add teleports.
for i, (x, y) in enumerate(self._teleport_info['teleport_in_xys']):
ET.SubElement(
worldbody,
'geom',
name=f'teleport_in_{i}',
type='cylinder',
size=f'{self._teleport_info["teleport_radius"]} .05',
pos=f'{x} {y} .05',
material='teleport_in',
contype='0',
conaffinity='0',
)
for i, (x, y) in enumerate(self._teleport_info['teleport_out_xys']):
ET.SubElement(
worldbody,
'geom',
name=f'teleport_out_{i}',
type='cylinder',
size=f'{self._teleport_info["teleport_radius"]} .05',
pos=f'{x} {y} .05',
material='teleport_out',
contype='0',
conaffinity='0',
)
if self._ob_type == 'pixels':
# Color wall.
wall = tree.find('.//material[@name="wall"]')
wall.set('rgba', '.6 .6 .6 1')
# Remove ambient light.
light = tree.find('.//light[@name="global"]')
light.attrib.pop('ambient')
# Remove torso light.
torso_light = tree.find('.//light[@name="torso_light"]')
torso_light_parent = tree.find('.//light[@name="torso_light"]/..')
torso_light_parent.remove(torso_light)
# Remove texture repeat.
grid = tree.find('.//material[@name="grid"]')
grid.set('texuniform', 'false')
if loco_env_type == 'ant':
# Color one leg white to break symmetry.
tree.find('.//geom[@name="aux_1_geom"]').set('material', 'self_white')
tree.find('.//geom[@name="left_leg_geom"]').set('material', 'self_white')
tree.find('.//geom[@name="left_ankle_geom"]').set('material', 'self_white')
else:
# Only show the target for states-based observation.
ET.SubElement(
worldbody,
'geom',
name='target',
type='cylinder',
size='.5 .05',
pos='0 0 .05',
material='target',
contype='0',
conaffinity='0',
)
def set_tasks(self):
# `tasks` is a list of tasks, where each task is a list of two tuples: (init_ij, goal_ij).
if self._maze_type == 'arena':
tasks = [
[(1, 1), (6, 6)],
]
elif self._maze_type == 'medium':
tasks = [
[(1, 1), (6, 6)],
[(6, 1), (1, 6)],
[(5, 3), (4, 2)],
[(6, 5), (6, 1)],
[(2, 6), (1, 1)],
]
elif self._maze_type == 'large':
tasks = [
[(1, 1), (7, 10)],
[(5, 4), (7, 1)],
[(7, 4), (1, 10)],
[(3, 8), (5, 4)],
[(1, 1), (5, 4)],
]
elif self._maze_type == 'giant':
tasks = [
[(1, 1), (10, 14)],
[(1, 14), (10, 1)],
[(8, 14), (1, 1)],
[(8, 3), (5, 12)],
[(5, 9), (3, 8)],
]
elif self._maze_type == 'teleport':
tasks = [
[(1, 10), (7, 1)],
[(1, 1), (7, 10)],
[(5, 6), (7, 10)],
[(7, 1), (7, 10)],
[(5, 6), (7, 1)],
]
else:
raise ValueError(f'Unknown maze type: {self._maze_type}')
self.task_infos = []
for i, task in enumerate(tasks):
self.task_infos.append(
dict(
task_name=f'task{i + 1}',
init_ij=task[0],
init_xy=self.ij_to_xy(task[0]),
goal_ij=task[1],
goal_xy=self.ij_to_xy(task[1]),
)
)
if self._reward_task_id == 0:
self._reward_task_id = 1 # Default task.
def initialize_renderer(self):
# Make custom renderer.
self.custom_renderer = mujoco.Renderer(
self.model,
width=self.width,
height=self.height,
)
self.render()
def reset(self, options=None, *args, **kwargs):
if options is None:
options = {}
# Set the task goal.
if self._reward_task_id is not None:
# Use the pre-defined task.
assert 1 <= self._reward_task_id <= self.num_tasks, f'Task ID must be in [1, {self.num_tasks}].'
self.cur_task_id = self._reward_task_id
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
elif 'task_id' in options:
# Use the pre-defined task.
assert 1 <= options['task_id'] <= self.num_tasks, f'Task ID must be in [1, {self.num_tasks}].'
self.cur_task_id = options['task_id']
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
elif 'task_info' in options:
# Use the provided task information.
self.cur_task_id = None
self.cur_task_info = options['task_info']
else:
# Randomly sample a task.
self.cur_task_id = np.random.randint(1, self.num_tasks + 1)
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
# Whether to provide a rendering of the goal.
render_goal = False
if 'render_goal' in options:
render_goal = options['render_goal']
# Get initial and goal positions with noise.
init_xy = self.add_noise(self.ij_to_xy(self.cur_task_info['init_ij']))
goal_xy = self.ij_to_xy(self.cur_task_info['goal_ij'])
if self._add_noise_to_goal:
goal_xy = self.add_noise(goal_xy)
# First, force set the position to the goal position to obtain the goal observation.
super().reset(*args, **kwargs)
# Do a few random steps to stabilize the environment.
num_random_actions = 40 if loco_env_type == 'humanoid' else 5
for _ in range(num_random_actions):
super().step(self.action_space.sample())
# Save the goal observation.
self.set_goal(goal_xy=goal_xy)
self.set_xy(goal_xy)
goal_ob = self.get_oracle_rep() if self._use_oracle_rep else self.get_ob()
if render_goal:
goal_rendered = self.render()
# Now, do the actual reset.
ob, info = super().reset(*args, **kwargs)
self.set_goal(goal_xy=goal_xy)
self.set_xy(init_xy)
ob = self.get_ob()
info['goal'] = goal_ob
if render_goal:
info['goal_rendered'] = goal_rendered
return ob, info
def step(self, action):
if self._success_timing == 'pre':
success = self.compute_success()
ob, reward, terminated, truncated, info = super().step(action)
if self._success_timing == 'post':
success = self.compute_success()
if self._teleport_info is not None:
# Check if the agent is close to a inbound teleport.
for x, y in self._teleport_info['teleport_in_xys']:
if np.linalg.norm(self.get_xy() - np.array([x, y])) <= self._teleport_info['teleport_radius'] * 1.5:
# Teleport the agent to a random outbound teleport.
teleport_out_xy = self._teleport_info['teleport_out_xys'][
np.random.randint(len(self._teleport_info['teleport_out_xys']))
]
self.set_xy(np.array(teleport_out_xy))
break
if success:
if self._terminate_at_goal:
terminated = True
info['success'] = 1.0
reward = 1.0
else:
info['success'] = 0.0
reward = 0.0
# If the environment is in the single-task mode, modify the reward.
if self._reward_task_id is not None:
reward = reward - 1.0 # -1 (failure) or 0 (success).
return ob, reward, terminated, truncated, info
def render(self):
if self.custom_renderer is None:
self.initialize_renderer()
self.custom_renderer.update_scene(self.data, camera=self.custom_camera)
return self.custom_renderer.render()
def get_ob(self, ob_type=None):
ob_type = self._ob_type if ob_type is None else ob_type
if ob_type == 'states':
return super().get_ob()
else:
frame = self.render()
return frame
def get_oracle_rep(self):
"""Return the oracle goal representation (i.e., the goal position)."""
return np.array(self.cur_goal_xy)
def compute_success(self):
if np.linalg.norm(self.get_xy() - self.cur_goal_xy) <= self._goal_tol:
return True
else:
return False
def set_goal(self, goal_ij=None, goal_xy=None):
"""Set the goal position and update the target object."""
if goal_xy is None:
self.cur_goal_xy = self.ij_to_xy(goal_ij)
if self._add_noise_to_goal:
self.cur_goal_xy = self.add_noise(self.cur_goal_xy)
else:
self.cur_goal_xy = goal_xy
if self._ob_type == 'states':
self.model.geom('target').pos[:2] = goal_xy
def get_oracle_subgoal(self, start_xy, goal_xy):
"""Get the oracle subgoal for the agent.
If the goal is unreachable, it returns the current position as the subgoal.
Args:
start_xy: Starting position of the agent.
goal_xy: Goal position of the agent.
Returns:
A tuple of the oracle subgoal and the BFS map.
"""
start_ij = self.xy_to_ij(start_xy)
goal_ij = self.xy_to_ij(goal_xy)
# Run BFS to find the next subgoal.
bfs_map = self.maze_map.copy()
for i in range(self.maze_map.shape[0]):
for j in range(self.maze_map.shape[1]):
bfs_map[i][j] = -1
bfs_map[goal_ij[0], goal_ij[1]] = 0
queue = [goal_ij]
while len(queue) > 0:
i, j = queue.pop(0)
for di, dj in [(-1, 0), (0, -1), (1, 0), (0, 1)]:
ni, nj = i + di, j + dj
if (
0 <= ni < self.maze_map.shape[0]
and 0 <= nj < self.maze_map.shape[1]
and self.maze_map[ni, nj] == 0
and bfs_map[ni, nj] == -1
):
bfs_map[ni][nj] = bfs_map[i][j] + 1
queue.append((ni, nj))
# Find the subgoal that attains the minimum BFS value.
subgoal_ij = start_ij
for di, dj in [(-1, 0), (0, -1), (1, 0), (0, 1)]:
ni, nj = start_ij[0] + di, start_ij[1] + dj
if (
0 <= ni < self.maze_map.shape[0]
and 0 <= nj < self.maze_map.shape[1]
and self.maze_map[ni, nj] == 0
and bfs_map[ni, nj] < bfs_map[subgoal_ij[0], subgoal_ij[1]]
):
subgoal_ij = (ni, nj)
subgoal_xy = self.ij_to_xy(subgoal_ij)
return np.array(subgoal_xy), bfs_map
def xy_to_ij(self, xy):
maze_unit = self._maze_unit
i = int((xy[1] + self._offset_y + 0.5 * maze_unit) / maze_unit)
j = int((xy[0] + self._offset_x + 0.5 * maze_unit) / maze_unit)
return i, j
def ij_to_xy(self, ij):
i, j = ij
x = j * self._maze_unit - self._offset_x
y = i * self._maze_unit - self._offset_y
return x, y
def add_noise(self, xy):
random_x = np.random.uniform(low=-self._noise, high=self._noise) * self._maze_unit / 4
random_y = np.random.uniform(low=-self._noise, high=self._noise) * self._maze_unit / 4
return xy[0] + random_x, xy[1] + random_y
class BallEnv(MazeEnv):
def update_tree(self, tree):
super().update_tree(tree)
# Add ball.
worldbody = tree.find('.//worldbody')
ball = ET.SubElement(worldbody, 'body', name='ball', pos='0 0 0.5')
ET.SubElement(ball, 'freejoint', name='ball_root')
ET.SubElement(
ball,
'geom',
name='ball',
size='.25',
material='ball',
priority='1',
conaffinity='1',
condim='6',
)
ET.SubElement(ball, 'light', name='ball_light', pos='0 0 4', mode='trackcom')
def set_tasks(self):
# `tasks` is a list of tasks, where each task is a list of three tuples: (agent_init_ij, ball_init_ij,
# goal_ij).
if self._maze_type == 'arena':
tasks = [
[(1, 6), (2, 3), (5, 2)],
[(2, 2), (5, 5), (2, 2)],
[(6, 1), (2, 3), (6, 6)],
[(6, 6), (1, 1), (6, 1)],
[(4, 6), (6, 2), (1, 6)],
]
elif self._maze_type == 'medium':
tasks = [
[(1, 1), (3, 4), (6, 6)],
[(6, 1), (6, 5), (1, 1)],
[(5, 3), (4, 2), (6, 5)],
[(6, 5), (1, 1), (5, 3)],
[(1, 6), (6, 1), (1, 6)],
]
else:
raise ValueError(f'Unknown maze type: {self._maze_type}')
self.task_infos = []
for i, task in enumerate(tasks):
self.task_infos.append(
dict(
task_name=f'task{i + 1}',
agent_init_ij=task[0],
agent_init_xy=self.ij_to_xy(task[0]),
ball_init_ij=task[1],
ball_init_xy=self.ij_to_xy(task[1]),
goal_ij=task[2],
goal_xy=self.ij_to_xy(task[2]),
)
)
if self._reward_task_id == 0:
self._reward_task_id = 4 # Default task.
def reset(self, options=None, *args, **kwargs):
if options is None:
options = {}
# Set the task goal.
if self._reward_task_id is not None:
# Use the pre-defined task.
assert 1 <= self._reward_task_id <= self.num_tasks, f'Task ID must be in [1, {self.num_tasks}].'
self.cur_task_id = self._reward_task_id
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
elif 'task_id' in options:
# Use the pre-defined task.
assert 1 <= options['task_id'] <= self.num_tasks, f'Task ID must be in [1, {self.num_tasks}].'
self.cur_task_id = options['task_id']
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
elif 'task_info' in options:
# Use the provided task information.
self.cur_task_id = None
self.cur_task_info = options['task_info']
else:
# Randomly sample a task.
self.cur_task_id = np.random.randint(1, self.num_tasks + 1)
self.cur_task_info = self.task_infos[self.cur_task_id - 1]
# Whether to provide a rendering of the goal.
render_goal = False
if 'render_goal' in options:
render_goal = options['render_goal']
# Get initial and goal positions with noise.
agent_init_xy = self.add_noise(self.ij_to_xy(self.cur_task_info['agent_init_ij']))
ball_init_xy = self.add_noise(self.ij_to_xy(self.cur_task_info['ball_init_ij']))
goal_xy = self.ij_to_xy(self.cur_task_info['goal_ij'])
if self._add_noise_to_goal:
goal_xy = self.add_noise(goal_xy)
# First, force set the position to the goal position to obtain the goal observation.
super(MazeEnv, self).reset(*args, **kwargs)
# Do a few random steps to stabilize the environment.
for _ in range(10):
super(MazeEnv, self).step(self.action_space.sample())
# Save the goal observation.
self.set_goal(goal_xy=goal_xy)
self.set_agent_ball_xy(goal_xy, goal_xy)
goal_ob = self.get_oracle_rep() if self._use_oracle_rep else self.get_ob()
if render_goal:
goal_rendered = self.render()
# Now, do the actual reset.
ob, info = super(MazeEnv, self).reset(*args, **kwargs)
self.set_goal(goal_xy=goal_xy)
self.set_agent_ball_xy(agent_init_xy, ball_init_xy)
ob = self.get_ob()
info['goal'] = goal_ob
if render_goal:
info['goal_rendered'] = goal_rendered
return ob, info
def step(self, action):
if self._success_timing == 'pre':
success = self.compute_success()
ob, reward, terminated, truncated, info = super(MazeEnv, self).step(action)
if self._success_timing == 'post':
success = self.compute_success()
if success:
if self._terminate_at_goal:
terminated = True
info['success'] = 1.0
reward = 1.0
else:
info['success'] = 0.0
reward = 0.0
# If the environment is in the single-task mode, modify the reward.
if self._reward_task_id is not None:
reward = reward - 1.0 # -1 (failure) or 0 (success).
return ob, reward, terminated, truncated, info
def compute_success(self):
if np.linalg.norm(self.get_agent_ball_xy()[1] - self.cur_goal_xy) <= self._goal_tol:
return True
else:
return False
def get_agent_ball_xy(self):
agent_xy = self.data.qpos[:2].copy()
ball_xy = self.data.qpos[-7:-5].copy()
return agent_xy, ball_xy
def set_agent_ball_xy(self, agent_xy, ball_xy):
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
qpos[:2] = agent_xy
qpos[-7:-5] = ball_xy
self.set_state(qpos, qvel)
if maze_env_type == 'maze':
return MazeEnv(*args, **kwargs)
elif maze_env_type == 'ball':
return BallEnv(*args, **kwargs)
else:
raise ValueError(f'Unknown maze environment type: {maze_env_type}')
================================================
FILE: ogbench/locomaze/point.py
================================================
import os
import gymnasium
import mujoco
import numpy as np
from gymnasium import utils
from gymnasium.envs.mujoco import MujocoEnv
from gymnasium.spaces import Box
class PointEnv(MujocoEnv, utils.EzPickle):
"""PointMass environment.
This is a simple 2-D point mass environment, where the agent is controlled by an x-y action vector that is added to
the current position of the point mass.
"""
xml_file = os.path.join(os.path.dirname(__file__), 'assets', 'point.xml')
metadata = {
'render_modes': ['human', 'rgb_array', 'depth_array'],
'render_fps': 10,
}
if gymnasium.__version__ >= '1.1.0':
metadata['render_modes'] += ['rgbd_tuple']
def __init__(
self,
xml_file=None,
render_mode='rgb_array',
width=200,
height=200,
**kwargs,
):
"""Initialize the Humanoid environment.
Args:
xml_file: Path to the XML description (optional).
render_mode: Rendering mode.
width: Width of the rendered image.
height: Height of the rendered image.
**kwargs: Additional keyword arguments.
"""
if xml_file is None:
xml_file = self.xml_file
utils.EzPickle.__init__(
self,
xml_file,
**kwargs,
)
observation_space = Box(low=-np.inf, high=np.inf, shape=(6,), dtype=np.float64)
MujocoEnv.__init__(
self,
xml_file,
frame_skip=5,
observation_space=observation_space,
render_mode=render_mode,
width=width,
height=height,
**kwargs,
)
def step(self, action):
prev_qpos = self.data.qpos.copy()
prev_qvel = self.data.qvel.copy()
action = 0.2 * action
self.data.qpos[:] = self.data.qpos + action
self.data.qvel[:] = np.array([0.0, 0.0])
mujoco.mj_step(self.model, self.data, nstep=self.frame_skip)
qpos = self.data.qpos.flat.copy()
qvel = self.data.qvel.flat.copy()
observation = self.get_ob()
if self.render_mode == 'human':
self.render()
return (
observation,
0.0,
False,
False,
{
'xy': self.get_xy(),
'prev_qpos': prev_qpos,
'prev_qvel': prev_qvel,
'qpos': qpos,
'qvel': qvel,
},
)
def get_ob(self):
return self.data.qpos.flat.copy()
def reset_model(self):
qpos = self.init_qpos + self.np_random.uniform(size=self.model.nq, low=-0.1, high=0.1)
qvel = self.init_qvel + self.np_random.standard_normal(self.model.nv) * 0.1
self.set_state(qpos, qvel)
return self.get_ob()
def get_xy(self):
return self.data.qpos.copy()
def set_xy(self, xy):
qpos = self.data.qpos.copy()
qvel = self.data.qvel.copy()
qpos[:] = xy
self.set_state(qpos, qvel)
================================================
FILE: ogbench/manipspace/__init__.py
================================================
from gymnasium.envs.registration import register
visual_dict = dict(
ob_type='pixels',
width=64,
height=64,
visualize_info=False,
)
cube_singletask_dict = dict(
permute_blocks=False,
success_timing='pre',
)
scene_singletask_dict = dict(
permute_blocks=False,
success_timing='pre',
)
puzzle_singletask_dict = dict(
success_timing='pre',
)
# Environments for offline goal-conditioned RL.
register(
id='cube-single-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=200,
kwargs=dict(
env_type='single',
),
)
register(
id='visual-cube-single-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=200,
kwargs=dict(
env_type='single',
**visual_dict,
),
)
register(
id='cube-double-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=500,
kwargs=dict(
env_type='double',
),
)
register(
id='visual-cube-double-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=500,
kwargs=dict(
env_type='double',
**visual_dict,
),
)
register(
id='cube-triple-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='triple',
),
)
register(
id='visual-cube-triple-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='triple',
**visual_dict,
),
)
register(
id='cube-quadruple-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='quadruple',
),
)
register(
id='visual-cube-quadruple-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='quadruple',
**visual_dict,
),
)
register(
id='cube-octuple-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1500,
kwargs=dict(
env_type='octuple',
),
)
register(
id='scene-v0',
entry_point='ogbench.manipspace.envs.scene_env:SceneEnv',
max_episode_steps=750,
kwargs=dict(
env_type='scene',
),
)
register(
id='visual-scene-v0',
entry_point='ogbench.manipspace.envs.scene_env:SceneEnv',
max_episode_steps=750,
kwargs=dict(
env_type='scene',
**visual_dict,
),
)
register(
id='puzzle-3x3-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='3x3',
),
)
register(
id='visual-puzzle-3x3-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='3x3',
**visual_dict,
),
)
register(
id='puzzle-4x4-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='4x4',
),
)
register(
id='visual-puzzle-4x4-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='4x4',
**visual_dict,
),
)
register(
id='puzzle-4x5-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x5',
),
)
register(
id='visual-puzzle-4x5-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x5',
**visual_dict,
),
)
register(
id='puzzle-4x6-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x6',
),
)
register(
id='visual-puzzle-4x6-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x6',
**visual_dict,
),
)
# Environments for reward-based single-task offline RL.
for task_id in [None, 1, 2, 3, 4, 5]:
task_suffix = '' if task_id is None else f'-task{task_id}'
reward_task_id = 0 if task_id is None else task_id # 0 means the default task.
register(
id=f'cube-single-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=200,
kwargs=dict(
env_type='single',
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'visual-cube-single-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=200,
kwargs=dict(
env_type='single',
**visual_dict,
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'cube-double-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=500,
kwargs=dict(
env_type='double',
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'visual-cube-double-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=500,
kwargs=dict(
env_type='double',
**visual_dict,
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'cube-triple-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='triple',
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'visual-cube-triple-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='triple',
**visual_dict,
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'cube-quadruple-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='quadruple',
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'visual-cube-quadruple-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='quadruple',
**visual_dict,
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'cube-octuple-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.cube_env:CubeEnv',
max_episode_steps=1500,
kwargs=dict(
env_type='octuple',
reward_task_id=reward_task_id,
**cube_singletask_dict,
),
)
register(
id=f'scene-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.scene_env:SceneEnv',
max_episode_steps=750,
kwargs=dict(
env_type='scene',
reward_task_id=reward_task_id,
**scene_singletask_dict,
),
)
register(
id=f'visual-scene-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.scene_env:SceneEnv',
max_episode_steps=750,
kwargs=dict(
env_type='scene',
**visual_dict,
reward_task_id=reward_task_id,
**scene_singletask_dict,
),
)
register(
id=f'puzzle-3x3-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='3x3',
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'visual-puzzle-3x3-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='3x3',
**visual_dict,
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'puzzle-4x4-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='4x4',
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'visual-puzzle-4x4-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=500,
kwargs=dict(
env_type='4x4',
**visual_dict,
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'puzzle-4x5-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x5',
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'visual-puzzle-4x5-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x5',
**visual_dict,
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'puzzle-4x6-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x6',
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
register(
id=f'visual-puzzle-4x6-singletask{task_suffix}-v0',
entry_point='ogbench.manipspace.envs.puzzle_env:PuzzleEnv',
max_episode_steps=1000,
kwargs=dict(
env_type='4x6',
**visual_dict,
reward_task_id=reward_task_id,
**puzzle_singletask_dict,
),
)
================================================
FILE: ogbench/manipspace/controllers/__init__.py
================================================
from ogbench.manipspace.controllers.diff_ik import DiffIKController
__all__ = ('DiffIKController',)
================================================
FILE: ogbench/manipspace/controllers/diff_ik.py
================================================
import mujoco
import numpy as np
PI = np.pi
PI_2 = 2 * np.pi
def angle_diff(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
return np.mod(q1 - q2 + PI, PI_2) - PI
class DiffIKController:
"""Differential inverse kinematics controller."""
def __init__(
self,
model: mujoco.MjModel,
sites: list,
qpos0: np.ndarray = None,
damping_coeff: float = 1e-12,
max_angle_change: float = np.radians(45),
):
self._model = model
self._data = mujoco.MjData(self._model)
self._qp0 = qpos0
self._max_angle_change = max_angle_change
# Cache references.
self._ns = len(sites) # Number of sites.
self._site_ids = np.asarray([self._model.site(s).id for s in sites])
# Preallocate arrays.
self._err = np.empty((self._ns, 6))
self._site_quat = np.empty((self._ns, 4))
self._site_quat_inv = np.empty((self._ns, 4))
self._err_quat = np.empty((self._ns, 4))
self._jac = np.empty((6 * self._ns, self._model.nv))
self._damping = damping_coeff * np.eye(6 * self._ns)
self._eye = np.eye(self._model.nv)
def _forward_kinematics(self) -> None:
"""Minimal computation required for forward kinematics."""
mujoco.mj_kinematics(self._model, self._data)
mujoco.mj_comPos(self._model, self._data) # Required for mj_jacSite.
def _integrate(self, update: np.ndarray) -> None:
"""Integrate the joint velocities in-place."""
mujoco.mj_integratePos(self._model, self._data.qpos, update, 1.0)
def _compute_translational_error(self, pos: np.ndarray) -> None:
"""Compute the error between the desired and current site positions."""
self._err[:, :3] = pos - self._data.site_xpos[self._site_ids]
def _compute_rotational_error(self, quat: np.ndarray) -> None:
"""Compute the error between the desired and current site orientations."""
for i, site_id in enumerate(self._site_ids):
mujoco.mju_mat2Quat(self._site_quat[i], self._data.site_xmat[site_id])
mujoco.mju_negQuat(self._site_quat_inv[i], self._site_quat[i])
mujoco.mju_mulQuat(self._err_quat[i], quat[i], self._site_quat_inv[i])
mujoco.mju_quat2Vel(self._err[i, 3:], self._err_quat[i], 1.0)
def _compute_jacobian(self) -> None:
"""Update site end-effector Jacobians."""
for i, site_id in enumerate(self._site_ids):
jacp = self._jac[6 * i : 6 * i + 3]
jacr = self._jac[6 * i + 3 : 6 * i + 6]
mujoco.mj_jacSite(self._model, self._data, jacp, jacr, site_id)
def _error_threshold_reached(self, pos_thresh: float, ori_thresh: float) -> bool:
"""Return True if position and rotation errors are below the thresholds."""
pos_achieved = np.linalg.norm(self._err[:, :3]) <= pos_thresh
ori_achieved = np.linalg.norm(self._err[:, 3:]) <= ori_thresh
return pos_achieved and ori_achieved
def _solve(self) -> np.ndarray:
"""Solve for joint velocities using damped least squares."""
H = self._jac @ self._jac.T + self._damping
x = self._jac.T @ np.linalg.solve(H, self._err.ravel())
if self._qp0 is not None:
jac_pinv = np.linalg.pinv(H)
q_err = angle_diff(self._qp0, self._data.qpos)
x += (self._eye - (self._jac.T @ jac_pinv) @ self._jac) @ q_err
return x
def _scale_update(self, update: np.ndarray) -> np.ndarray:
"""Scale down update so that the max allowable angle change is not exceeded."""
update_max = np.max(np.abs(update))
if update_max > self._max_angle_change:
update *= self._max_angle_change / update_max
return update
def solve(
self,
pos: np.ndarray,
quat: np.ndarray,
curr_qpos: np.ndarray,
max_iters: int = 20,
pos_thresh: float = 1e-4,
ori_thresh: float = 1e-4,
) -> np.ndarray:
self._data.qpos = curr_qpos
for _ in range(max_iters):
self._forward_kinematics()
self._compute_translational_error(np.atleast_2d(pos))
self._compute_rotational_error(np.atleast_2d(quat))
if self._error_threshold_reached(pos_thresh, ori_thresh):
break
self._compute_jacobian()
update = self._scale_update(self._solve())
self._integrate(update)
return self._data.qpos.copy()
================================================
FILE: ogbench/manipspace/descriptions/button_inner.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/button_outer.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/buttons.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/cube.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/cube_inner.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/cube_outer.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/drawer.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/floor_wall.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/robotiq_2f85/2f85.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/robotiq_2f85/LICENSE
================================================
Copyright (c) 2013, ROS-Industrial
All rights reserved.
Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:
Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
Redistributions in binary form must reproduce the above copyright notice, this
list of conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
================================================
FILE: ogbench/manipspace/descriptions/robotiq_2f85/README.md
================================================
# Robotiq 2F-85 Description (MJCF)
Requires MuJoCo 2.2.2 or later.
## Overview
This package contains a simplified robot description (MJCF) of the [Robotiq 85mm
2-Finger Adaptive
Gripper](https://robotiq.com/products/2f85-140-adaptive-robot-gripper) developed
by [Robotiq](https://robotiq.com/). It is derived from the [publicly available
URDF
description](https://github.com/ros-industrial/robotiq/tree/kinetic-devel/robotiq_2f_85_gripper_visualization).
## URDF → MJCF derivation steps
1. Added `` to the URDF's
`` clause in order to preserve visual geometries.
2. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
3. Manually edited the MJCF to extract common properties into the `` section.
4. Added `` clauses to prevent collisions between the linkage bodies.
5. Broke up collision pads into two pads for more contacts.
6. Increased pad friction and priority.
7. Added `impratio=10` for better noslip.
8. Added `scene.xml` which includes the robot, with a textured groundplane, skybox, and haze.
9. Added hanging box to `scene.xml`.
## License
This model is released under a [BSD-2-Clause License](LICENSE).
================================================
FILE: ogbench/manipspace/descriptions/robotiq_2f85/scene.xml
================================================
================================================
FILE: ogbench/manipspace/descriptions/universal_robots_ur5e/LICENSE
================================================
Copyright 2018 ROS Industrial Consortium
Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its contributors
may be used to endorse or promote products derived from this software without
specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
================================================
FILE: ogbench/manipspace/descriptions/universal_robots_ur5e/README.md
================================================
# Universal Robots UR5e Description (MJCF)
Requires MuJoCo 2.3.3 or later.
## Overview
This package contains a simplified robot description (MJCF) of the
[UR5e](https://www.universal-robots.com/products/ur5-robot/) developed by
[Universal Robots](https://www.universal-robots.com/). It is derived from the
[publicly available URDF
description](https://github.com/ros-industrial/universal_robot/tree/kinetic-devel/ur_e_description).
### URDF → MJCF derivation steps
1. Converted the DAE [mesh
files](https://github.com/ros-industrial/universal_robot/tree/kinetic-devel/ur_e_description/meshes/ur5e/visual)
to OBJ format using [Blender](https://www.blender.org/).
2. Processed `.obj` files with [`obj2mjcf`](https://github.com/kevinzakka/obj2mjcf).
3. Added `` to the URDF's
`` clause in order to preserve visual geometries.
4. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
5. Added a tracking light to the base.
6. Manually edited the MJCF to extract common properties into the `
