WAX-ML: A Python library for machine learning and feedback loops on streaming data

TL;DR

WAX-ML is a Python library that extends JAX to facilitate the development of machine learning algorithms and feedback loops on streaming time series data, integrating with pandas, xarray, and Gym.

Contribution

It introduces a comprehensive Python toolkit for designing online learning and reinforcement learning algorithms on streaming data, enhancing JAX with user-friendly features.

Findings

Enables easy implementation of feedback loops in streaming data scenarios.

Supports online and reinforcement learning algorithms with seamless integration.

Provides tools compatible with pandas, xarray, and Gym for end-user accessibility.

Abstract

Wax is what you put on a surfboard to avoid slipping. It is an essential tool to go surfing... We introduce WAX-ML a research-oriented Python library providing tools to design powerful machine learning algorithms and feedback loops working on streaming data. It strives to complement JAX with tools dedicated to time series. WAX-ML makes JAX-based programs easy to use for end-users working with pandas and xarray for data manipulation. It provides a simple mechanism for implementing feedback loops, allows the implementation of online learning and reinforcement learning algorithms with functions, and makes them easy to integrate by end-users working with the object-oriented reinforcement learning framework from the Gym library. It is released with an Apache open-source license on GitHub at https://github.com/eserie/wax-ml.