AI research startup DeepMind has now open-sourced new libraries for neural networks and reinforcement learning based on JAX. JAX is a numerical computing library launched by Google a couple of years ago, and can automatically differentiate native Python and NumPy functions. JAX uses XLA (Accelerated Linear Algebra) to compile and run your NumPy programs on GPUs and TPUs, which is great for machine learning to compute. 

Known as Haiku and RLax, both open-sourced libraries can be used by AI enthusiasts, professionals and researchers for reinforcement learning initiatives.

Haiku: Sonnet for JAX 

This Haiku library allows AI researchers to fully access JAX’s pure function transformations by providing object-oriented programming models. Haiku is a neural network library for JAX, meaning it is created on the programming model and APIs of Sonnet– a library on top of TensorFlow designed to provide simple, composable abstractions for machine learning research. Haiku, in fact, has been developed by some of the authors of Sonnet. The open-source library Sonnet is widely adopted at DeepMind.

Haiku has been experimented by researchers at DeepMind at scale with success. Haiku has demonstrated promising results in big-scale experimentation on image and language datasets using reinforcement learning. But, Haiku isn’t the only neural network library made for JAX, and others also exist such as Google’ Flax.

RLax, on the other hand, is a simple RL library for JAX. Instead of providing complete algorithms for RL, this library provides useful building blocks for implementing specific mathematical operations to build fully-functional RL agents. DeepMind pioneered deep reinforcement learning – to create the first artificial agents to achieve human-level performance across many challenging domains. Two years ago, the company rolled out the first widely successful algorithm for deep reinforcement learning. 

Advantages Of New DeepMind Open-Source Libraries

So, if there are many neural network libraries for JAX, why should you choose Haiku and RLax?

According to Deep Mind, Haiku has been created to make specific processes straightforward for managing things like model parameters along with the model state.

Haiku is presently in alpha, many researchers have experimented with Haiku for many months and have reproduced a number of experiments at a large scale.

DeepMind has already reproduced many experiments in Haiku and JAX, and include big-scale tests in image and language processing models, generative models, as well as reinforcement learning.

The advantage is that both libraries have well-defined instructions, and for Haiku, the APIs are similar to Sonnet, which means moving to Haiku is likely to be quite easy for those working on Sonnet. The fact that Haiku and RLax are supported and tested by Deep Mind can give developers, researchers and students peace of mind for their advanced ML endeavours.

Overview

It seems while project Haiku and RLax have been launched to GitHub, they are still experimental. While JAX had been created for high-performance machine learning research by giving an integrated system for learning model optimisation, it is not the most comfortable tool to operate. This is why experts think DeepMind developed Haiku and RLax.

Also Read: 10 Best ML Engineering Practices