Train a neural network to predict robot dynamics

Train a neural network in python to predict robot dynamics

Understand how to work with real robotics data.

My life as a robotics learning researcher is mostly data science (understanding data) and a bit of fun math (reinforcement learning). Here’s a snippet you can run for yourself and take lessons away into your projects (and career)!

Deep learning and robotics

Neural networks have become state-of-the-art in multiple digital domains, particularly computer vision and language processing. The neural networks are complex function approximations (capable of fitting any function sufficiently accurately, from a fun theorem ), so why not use them for robotics?

In robots, we want to model the dynamics for two things:

1. Planning to take action . Actions accomplish tasks. This makes useful robots. We all want useful robots.
2. Understanding the system . Understanding the system can open up more future actions and tasks to the robot.

The classic neural network illustration. Densely-connected, nonlinear activation functions can fit any function with a sufficient number of neurons.

The optimization problem

When learning with a neural network will predict a discrete step in the dynamics of the system. Dynamic systems take the form shown below:

• s is the state of the system (e.g. physical position).
• a is the action of the agent (e.g. motor voltage).
• f is the true dynamics of the robot.
• w is the disturbance (randomness, or un-modeled part of the system).

General form of a dynamic system.

What we will model with the neural network (intentionally) is the underlying function, f , and we implicitly capture the disturbance, w , in the data. How do we do this? For data “niceness” we do a couple of things. Niceness means that the data is close to uniformly distributed with a mean of 0 and a standard deviation of 1.

We model the learned dynamics (denoted by a theta subscript) to predict a change in state, rather than a true state. This delta state, to the right, is simply the change from the last time step to the next.

What does this look like in the context of a neural network? Specifically, we can minimize the norm squared of the difference between the change in state and the function approximator. This is a mean squared error . It’s implemented in PyTorch here .

State of the art is using things like the Adam optimizer , Swish activation function (smooth version of ReLU), a learning rate scheduler , and more. The important way to get NN’s to work is data normalization. If you look closely at the code, I handle it in a very modular way. Without it, fitting models to real data would be near impossible.

The code (PyTorch)

This is how you will run the code . I include a conda environment called ml (PyTorch, AI Gym, Mujoco, and more). The trainer code will be found in the learn directory. Parameters are stored in a configuration file.

(ml) user: dynamicslearn nato$ python learn/trainer.py robot=iono_sim

The robots

There are pre-filtered data in the repository from real flights of experimental robots learning control with reinforcement learning. It is an active research repository at University of California, Berkeley.

The code includes simulators for two flying robots, the Crazyflie and the Ionocraft . Both of these allow you to control the robots at 100 Hz, and the simulator calculates dynamics at 1000Hz. You can find the dynamic simulator here . The specific robot defines the force transforms for different actuators.

Rich NN Object

When working with a complex task (figuring out how to model a robot), it helps to have a configurable network object. You’ll see a series of flags here:

• A prob flag — this changes from MSE to an advanced loss function with better variance regularization from this paper .
• A history parameter — we want to pass past states into the predictor (when states evolve faster then some of the underlying dynamics evolve).
• A bunch of parameters — hidden width, depth, dropout, etc.
• SciKitLearn normalization scalars — super important for getting a workable optimization surface (you may get NaN’s without this).

Notice that I use X for states and U for actions.

class GeneralNN(nn.Module):
    def __init__(self, **nn_params):
        super(GeneralNN, self).__init__()
        # Store the parameters:
        self.prob = nn_params['training']['probl']
        self.hidden_w = nn_params['training']['hid_width']
        self.depth = nn_params['training']['hid_depth']        self.hist = nn_params['history']
        self.n_in_input = nn_params['du'] * (self.hist + 1) 
        self.n_in_state = nn_params['dx'] * (self.hist + 1)
        self.n_in = self.n_in_input + self.n_in_state
        self.n_out = nn_params['dt']        self.activation = Swish()  
        self.d = nn_params['training']['dropout']
        self.E = 0  # clarify that these models are not ensemble        self.scalarX = StandardScaler()  
        self.scalarU = MinMaxScaler(feature_range=(-1, 1))
        self.scalardX = MinMaxScaler(feature_range=(-1, 1))        self.init_training = False

Create layers dynamically

This may be the most useful part of the article for someone new to PyTorch.It lets you create a network of various depths and widths (the alternate is hard coding the configurations).

        layers = []
        layers.append(('dynm_input_lin', nn.Linear(
            self.n_in, self.hidden_w)))  # input layer
        layers.append(('dynm_input_act', self.activation))
        for d in range(self.depth):
            layers.append(
                ('d_l' + str(d), nn.Linear(self.hid_w, self.hid_w)))
            layers.append(('dynm_act_' + str(d), self.activation))
        layers.append(('d_o_l', nn.Linear(self.hid_w, self.n_out)))
        self.features = nn.Sequential(OrderedDict([*layers]))

Train the model

Omitting the optimizer code that is publicly available, and detailed in many tutorials , we can have a nice training function as our top-level object. Let lower levels of abstraction handle the details.

def train_model(X, U, dX, model_cfg, logged=False):
    if logged: log.info(f"Training Model on {np.shape(X)[0]} pts")
    start = time.time()
    train_log = dict()    train_log['model_params'] = model_cfg.params
    model = hydra.utils.instantiate(model_cfg)    model.train_cust((X_t, U_t, dX_t), model_cfg.params)    end = time.time()
    if logged: log.info(f"Trained Model in {end-start} s")
    return model, train_log

The neural network code is here . The training script is here . The overall code is here .