Intro

This tutorial is designed for those who have ventured into the domain of reinforcement learning using Isaac Gym Preview 4. You've spent hours meticulously training your models, but now, you face the pivotal challenge of transitioning these trained policies to run in different environments—be it for further testing, real-world applications, or integrating with other simulation tools. We will go through the process of converting your trained models from Isaac Gym to the ONNX format, ensuring they are ready for deployment in diverse environments or simulation frameworks.

I will be using and slightly modifying the Anymal Terrain environment for this example. We will use it to train a control policy that we will later export as an ONNX model. We will then perform inference with this model during the testing phase.

Before we get started

  1. Isaac Gym Preview : https://github.com/NVIDIA-Omniverse/IsaacGymEnvs

    1. This has been installed following the directions from the repo.

  2. RL Games (v1.6.1) : https://github.com/Denys88/rl_games

    1. Installed as a dependency of Isaac Gym

    2. We will replace the one installed with Isaac with our modified version later

  3. pip install

    1. onnx

    2. onnxruntime

  4. I put the fully modified files here just in case : https://github.com/Tbarkin121/IsaacExamples

Modifying RL Games for ONNX Export


RL Games is a framework for training the reinforcement learning models that is used by Isaac Gym. To export these models for broader use, specifically in ONNX format, we need to make a few adjustments within RL Games. In order to do this, we need to clone a copy of RL Games from the repo : https://github.com/Denys88/rl_games

  1. git clone https://github.com/Denys88/rl_games.git

  2. cd rl_games

  3. pip install -e .

    • this installs the RL Games repo in editable mode. We can make now make changes to the files in this directory and when we run train.py, they will be run.

The file we need to edit is torch_runner.py

  1. Add the ModelWrapper class (I threw mine in at the very bottom of torch_runner.py):

class ModelWrapper(torch.nn.Module):
    '''
    Main idea is to ignore outputs which we don't need from model
    '''
    def __init__(self, model):
        torch.nn.Module.__init__(self)
        self._model = model
        
        
    def forward(self,input_dict):
        input_dict['obs'] = self._model.norm_obs(input_dict['obs'])
        '''
        just model export doesn't work. Looks like onnx issue with torch distributions
        thats why we are exporting only neural network
        '''
        #print(input_dict)
        #output_dict = self._model.a2c_network(input_dict)
        #input_dict['is_train'] = False
        #return output_dict['logits'], output_dict['values']
        return self._model.a2c_network(input_dict)

2. Modify the run_play function

    def run_play(self, args):
        print('Started to play')
        player = self.create_player()
        _restore(player, args)
        _override_sigma(player, args)

        import rl_games.algos_torch.flatten as flatten
        inputs = {
            'obs' : torch.zeros((1,) + player.obs_shape).to(player.device),
            'rnn_states' : player.states,
        }

        with torch.no_grad():
            adapter = flatten.TracingAdapter(ModelWrapper(player.model), inputs, allow_non_tensor=True)
            traced = torch.jit.trace(adapter, adapter.flattened_inputs, check_trace=False)
            flattened_outputs = traced(*adapter.flattened_inputs)
            print(flattened_outputs)

        torch.onnx.export(traced, *adapter.flattened_inputs, player.config['name']+".onnx", verbose=True, input_names=['obs'], output_names=['mu','log_std', 'value'])

        player.run()

When you run test=True, the play_run function is used. An onnx file with the name of the task will be saved in the directory you are running train.py from.

Testing the ONNX Model : Light Software-in-the-Loop

To show that the onnx model we just generated is working correctly, we can modify anymal_terrain.py to feed inferences from the onnx model into env0.

Work in Progress

Additional Resources

https://github.com/Denys88/rl_games/blob/master/notebooks/train_and_export_onnx_example_continuous.ipynb

Really it took me longer than I’d like to admit to understand that RL Games had an example of how to export onnx files. I knew it existed for a while I wasn’t sure how to apply this to training with Isaac.