This project focuses on training a Unitree Go2 quadruped robot using Reinforcement Learning. It utilizes the PPO (Proximal Policy Optimization) algorithm, implemented from scratch using JAX, Flax, and Optax, to train the robot for various tasks.

The robot models are included via the mujoco_menagerie Git submodule.

• ppo_agent.py: Defines the core PPO agent. This includes the ActorCritic neural network model (built with Flax) and all the necessary logic for action selection, loss calculation, and model updates using JAX and Optax.

• train_standing.py: A complete training script to teach the Go2 robot to stand still. It features a custom Gymnasium environment (Go2StandEnv) with a specific reward function designed to incentivize maintaining a target height while minimizing body and joint velocity. Running this script will start the training process and visualize the agent's performance at the end.

• spawn_doggo.py: A simple utility script to visualize the Go2 robot model within its MuJoCo environment. It loads the scene from the mujoco_menagerie submodule and has the robot perform random actions. This is useful for verifying the environment setup and ensuring the model loads and renders correctly.

• go2_train.py: A more general training script that uses the PPO agent to train a policy for a pre-registered Gymnasium environment named "Go2Gait-v0". This serves as a template for training other behaviors.

• requirements.txt: Lists the necessary Python packages required to set up the environment and run the project.

• mujoco_menagerie/: A Git submodule containing the collection of MuJoCo models, including the Unitree Go2 model used in the simulation environments.

1. Clone the repository:

   git clone <your-repo-url>
   cd RL_doggo

2. Initialize and update the submodule: This step is crucial as it will download the mujoco_menagerie models.

   git submodule update --init --recursive

3. Create and activate a virtual environment:

   python3 -m venv rlenv
   source rlenv/bin/activate

4. Install the required packages:

   pip install -r requirements.txt

Make sure you have activated the virtual environment (source rlenv/bin/activate) before running any of the scripts.

• To visualize the robot model performing random actions:

  python spawn_doggo.py

• To train the robot to stand:

  python train_standing.py