JaxGCRL

JaxGCRL provides 8 fast, goal-conditioned environments built on MJX and BRAX. These environments are vectorized and JIT-compiled with JAX for rapid experimentation in self-supervised reinforcement learning.

| Environment | Env name | 
| :- | :-: |
| Reacher | `reacher` |
| Half Cheetah | `cheetah` |
| Pusher | `pusher_easy` <br> `pusher_hard` |
| Ant | `ant` |
| Ant Maze | `ant_u_maze` <br> `ant_big_maze` <br> `ant_hardest_maze` |
| Ant Soccer | `ant_ball` |
| Ant Push | `ant_push` |
| Humanoid | `humanoid` |

This package provides a simplified MuJoCo (MJCF) robot description of the Franka Emika Panda. It is derived from the official URDF description and is optimized for use in MuJoCo-based environments.

Requirements:

• MuJoCo 2.3.3 or later.

JaxGCRL supports four categories of manipulation tasks. The robot arm is 7DoF (with a gripper if relevant), but only 4 joints are controllable.

Action Space Note: The action space uses absolute joint angle control. This is notably different from environments using absolute or delta end-effector (EEF) position control and may increase task difficulty.

Supported categories:

• Reach
• Grasp
• Push (includes easy and hard variants)
• Binpick (includes easy, hard, and easy EEF variants)

EEF Variant: The binpick easy EEF variant uses only the end-effector hand and rectangular fingers, utilizing delta end-effector position control instead of joint angle control.

To use these environments, specify the environment name. It is recommended to adjust the episode length to the value specified in the environment's file (self.episode_length).

The core algorithm in this repository is Contrastive Reinforcement Learning (CRL). It learns a critic function without explicit rewards by contrasting positive samples (states sampled from the future of a given trajectory) against negative samples (states sampled from a random trajectory).

Key configurable components of the CRL implementation include:

• Energy Functions: Determines how state similarity is measured. Supported options are norm, l2, dot, and cosine.
• Contrastive Losses: The loss function used for training. Supported options are binary_nce, fwd_infonce, bwd_infonce, and sym_infonce.
• Regularization: The logsumexp_penalty can be applied to the contrastive loss to improve performance and is mathematically required for the fwd_infonce loss.

JaxGCRL integrates with Weights & Biases (Wandb) for tracking metrics and rendering final policies as artifacts. To enable logging, use the --log_wandb flag.

You can organize your experiments using these flags:

• --wandb_project_name: The name of the Wandb project.
• --wandb_group: The group name within the project.
• --exp_name: The specific name for the experiment.

To run a hyperparameter sweep, use the Wandb CLI to initialize the sweep from a YAML config and then run the agent.

# Example: Running a loop of experiments with Wandb logging
for c_loss in binary infonce flatnce fb dp; do
	for seed in 1 2 3 4 5; do
		python run.py crl --seed ${seed} \
		--wandb_project_name crl --wandb_group contrastive_loss_experiments \ 
		--exp_name ${c_loss} \ 
		--contrastive_loss_fn ${c_loss} --energy_fn norm \ 
        --log_wandb
	done
done

To start a basic Contrastive Reinforcement Learning (CRL) experiment, use the run.py script with the crl command and specify an environment using the --env flag. For example, to run on the ant environment:

python run.py crl --env ant

To register and use a new contrastive loss function (e.g., your_loss) for the CRL algorithm, follow these steps:

1. Register the loss: Open agents/crl/losses.py and add your new function to the contrastive_loss_fn function:

   if name == "your_loss":
       critic_loss = ... # your implementation

2. Run training: Use the --contrastive_loss_fn flag in your execution command to specify the new loss.

python run.py crl --contrastive_loss_fn "your_loss"

To verify your installation, run a quick test experiment using the jaxgcrl command. The jaxgcrl command is a wrapper around python run.py.

jaxgcrl crl --env ant

CLI Usage and Flags

• Use jaxgcrl --help to see available agents.
• Use jaxgcrl {crl,ppo,sac,td3} --help to see available configurations for a specific agent.

Common Flags:

• --env=...: Specify the environment name (e.g., ant). See jaxgcrl/utils/env.py for the full list.
• --log_wandb: Enables Wandb logging. Remove this flag to omit logging.
• --total_env_steps: Set the duration of the run.
• --num_envs: Set the number of environments (adjust based on GPU memory).
• Algorithmic flags: --contrastive_loss_fn, --energy_fn, --h_dim, --n_hidden, etc.

To integrate a custom model architecture into the CRL algorithm, you must modify agents/crl/networks.py:

1. Critic: Extend or modify the MLP class, or create a new model specifically designed for generating contrastive embeddings.
2. Actor: Implement and provide an appropriate make_policy_network function that defines your desired actor architecture.

JaxGCRL has built-in support for Weights & Biases (Wandb) to track metrics and visualize learned policies as artifacts.

Enabling Wandb

Enable logging by adding the --log_wandb flag to your command. You can also organize your experiments using these flags:

• --project_name
• --group_name
• --exp_name

Running Wandb Sweeps

To run a hyperparameter sweep, use the following workflow:

1. Start the sweep:

   wandb sweep --project example_sweep ./scripts/sweep.yml

2. Run the agent using the output from the previous command:

   wandb agent <previous_command_output>

In addition to CRL, the repository provides several baseline RL algorithms based on Brax implementations:

• PPO
• SAC
• TD3

Configuration Options

• Hindsight Experience Replay (HER): For SAC and TD3, you can enable the --use_her flag to utilize the Hindsight Experience Replay buffer, which is beneficial in sparse reward settings.
• Reward Types: While goal-conditioned RL often doesn't assume rewards, most environments support both sparse and dense rewards. You can switch to dense rewards using the --use_dense_reward flag.

# Example: Running SAC with HER and dense rewards
python run.py sac --use_her --use_dense_reward