What Are DeepMind’s Newly Released Libraries For Neural Networks & Reinforcement Learning?

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.

Subscribe to our Newsletter

Join our editors every weekday evening as they steer you through the most significant news of the day, introduce you to fresh perspectives, and provide unexpected moments of joy
Your newsletter subscriptions are subject to AIM Privacy Policy and Terms and Conditions.

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.


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

Vishal Chawla
Vishal Chawla is a senior tech journalist at Analytics India Magazine and writes about AI, data analytics, cybersecurity, cloud computing, and blockchain. Vishal also hosts AIM's video podcast called Simulated Reality- featuring tech leaders, AI experts, and innovative startups of India.

Download our Mobile App

MachineHack | AI Hackathons, Coding & Learning

Host Hackathons & Recruit Great Data Talent!

AIMResearch Pioneering advanced AI market research

With a decade of experience under our belt, we are transforming how businesses use AI & data-driven insights to succeed.

The Gold Standard for Recognizing Excellence in Data Science and Tech Workplaces

With Best Firm Certification, you can effortlessly delve into the minds of your employees, unveil invaluable perspectives, and gain distinguished acclaim for fostering an exceptional company culture.

AIM Leaders Council

World’s Biggest Community Exclusively For Senior Executives In Data Science And Analytics.

3 Ways to Join our Community

Telegram group

Discover special offers, top stories, upcoming events, and more.

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Subscribe to our Daily newsletter

Get our daily awesome stories & videos in your inbox