34. Reinforcement Learning

This section shows how to use AGX dynamics with Gymnasium-style environments to train reinforcement learning (RL) control policies.

../_images/RL-loop.png

You will learn how to:

  • Wrap an AGX simulation as a Gymnasium environment

  • Run and visualize example environments

  • Train controllers using popular RL libraries

34.1. Installation

Install required Python packages:

pip install -r data/python/RL/requirements.txt

Note

We do not support platforms with python versions < 3.7.

34.2. AGX gym environment

An AGXGymEnv inherits from Gymnasium.Env and provides:

  • Implementations of the API methods step() render(), close(), reset()

  • Custom hooks for:

    • _build_scene() - Sets up the AGX dynamics simulation (environment).

    • _set_action() - Passes the action to the controller. Called in step().

    • _observe() - Returns (obs, reward, terminal, truncated, info). Called in env.step().

  • Optional camera rendering and graphics.

34.3. Training a control policy

These examples use an AGX implementation of the classic reinforcement learning problem called cart pole, where the goal is to balance a pole on a moving cart. As opposed to classic CartPole, our environment uses continuous actions and the observation space is either the standard state, or visual input from a virtual camera. We show how to train a control policy using stable-baselines3 and run trained control policies.

Run the CartPole environment with a random policy:

python data/python/RL/cartpole.py

Train a control policy from scratch:

python data/python/RL/cartpole.py --train

Train using visual observation:

python data/python/RL/cartpole.py --train --observation-space visual

Train a policy on 4 parallel environments:

python data/python/RL/cartpole.py --train --num-env 4 --observation-space visual

Load a control policy:

python data/python/RL/cartpole.py --load path/to/trained/policy

Load one of our pre-trained policies:

python data/python/RL/cartpole.py --observation-space visual --load data/python/RL/policyModels/cartpole_visual_policy.zip