CINXE.COM

<!DOCTYPE html> <html class='v2' dir='ltr' lang='en' xmlns='http://www.w3.org/1999/xhtml' xmlns:b='http://www.google.com/2005/gml/b' xmlns:data='http://www.google.com/2005/gml/data' xmlns:expr='http://www.google.com/2005/gml/expr'> <head> <link href='https://www.blogger.com/static/v1/widgets/3566091532-css_bundle_v2.css' rel='stylesheet' type='text/css'/> <meta content='text/html; charset=UTF-8' http-equiv='Content-Type'/> <meta content='blogger' name='generator'/> <link href='https://blog.tensorflow.org/favicon.ico' rel='icon' type='image/x-icon'/> <link href='https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html' rel='canonical'/> <link rel="alternate" type="application/atom+xml" title="The TensorFlow Blog - Atom" href="https://blog.tensorflow.org/feeds/posts/default" /> <link rel="alternate" type="application/rss+xml" title="The TensorFlow Blog - RSS" href="https://blog.tensorflow.org/feeds/posts/default?alt=rss" /> <link rel="service.post" type="application/atom+xml" title="The TensorFlow Blog - Atom" href="https://www.blogger.com/feeds/7864883956188652345/posts/default" /> <link rel="alternate" type="application/atom+xml" title="The TensorFlow Blog - Atom" href="https://blog.tensorflow.org/feeds/6024940900969663155/comments/default" />  <link href='https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png' rel='image_src'/> <meta content='In this blog post, we will show you how to train a game agent using reinforcement learning using JAX/Flax, convert the model to TensorFlow Lite, and d' name='description'/> <meta content='https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html' property='og:url'/> <meta content='Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ' property='og:title'/> <meta content='In this blog post, we will show you how to train a game agent using reinforcement learning using JAX/Flax, convert the model to TensorFlow Lite, and d' property='og:description'/> <meta content='https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/w1200-h630-p-k-no-nu/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png' property='og:image'/> <meta charset='UTF-8'/> <meta content='IE=edge' http-equiv='X-UA-Compatible'/> <meta content='width=device-width, initial-scale=1' name='viewport'/> <meta content='https://www.gstatic.com/tf_blog/images/image_blank.png' property='og:image'/> <meta content='https://www.gstatic.com/tf_blog/images/image_blank.png' property='twitter:image'/> <meta content='summary_large_image' name='twitter:card'/> <meta content='Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ' property='twitter:title'/> <title>Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite — The TensorFlow Blog</title> <style id='page-skin-1' type='text/css'></style>  <link href='https://fonts.googleapis.com/css?family=Google+Sans:400,500,700|Roboto:400,400italic,500,500italic,700,700italic|Roboto+Mono:400,500,700|Material+Icons' rel='stylesheet'/>   <link href="//prismjs.com/themes/prism.css" rel="stylesheet"> <script src="//prismjs.com/prism.js" type="text/javascript"></script> <script src='https://cdnjs.cloudflare.com/ajax/libs/prism/1.17.1/plugins/autoloader/prism-autoloader.min.js'></script>   <link href='https://cdn.jsdelivr.net/npm/zoom-vanilla.js/dist/zoom.css' rel='stylesheet'/> <script defer='defer' src='https://cdn.jsdelivr.net/npm/zoom-vanilla.js/dist/zoom-vanilla.min.js' type='text/javascript'></script>  <link href='https://www.gstatic.com/tf_blog/images/favicon.png' rel='shortcut icon' type='image/png'/> <script type='text/javascript'> //<![CDATA[ const qs = (string, el = document) => el.querySelector(string); const qsa = (string, el = document) => el.querySelectorAll(string); class App { constructor() { this.body = qs('body'); this.detailBody = qs('.tensorsite-detail__body'); this.overlay = qs('.header__overlay'); this.hamburger = qs('.header__hamburger'); this.sideMenu = qs('.header__side-menu'); this.detailBodies = qsa('.tensorsite-detail__body'); this.searchForms = qsa('.searchbox'); this.searchInputs = qsa('.search-input'); this.homeHref = qs('#home-href'); this.featuredCard = qs('.tensorsite-card.featured'); this.featuredPostHref = this.featuredCard && this.featuredCard .querySelector('.tensorsite-card__href') .getAttribute('href'); this.cards = qsa('.tensorsite-card'); this.images = qsa('img[border]'); this.cardDescriptions = qsa('.tensorsite-content__description'); this.hiddenDescription = qsa('.tensorsite-detail__description'); this.iconLinks = qs('.social-icons__links').children this.iconTooltips = qsa('[class^="icon-tooltip"]') this._toggleMobileMenu = this._toggleMobileMenu.bind(this); this._closeMenu = this._closeMenu.bind(this); this._onResize = this._onResize.bind(this); this._getScreen = this._getScreen.bind(this); this._searchGoogle = this._searchGoogle.bind(this); this._handleSearchKeypress = this._handleSearchKeypress.bind(this); this._removeDividerAboveImage(); this._setAllTagActive(); this._showFeaturedPost(); this._redirectWithMaxResults(); this._makeImagesZoomable(); this._removeCardLineBreaks(); this._getNextPost().then(()=>{ this._removeCardLineBreaks(); }) this.addEventListeners(); } addEventListeners() { window.addEventListener('resize', this._onResize); this.hamburger.addEventListener('click', this._toggleMobileMenu); this.searchForms.forEach(el => el.addEventListener('submit', this._searchGoogle)); this.searchInputs.forEach(el => el.addEventListener('keypress', this._handleSearchKeypress)); Array.from(this.iconLinks).forEach((icon, i) => { icon.addEventListener("mouseover", () => icon.querySelectorAll('[class^="icon-tooltip"]')[0].style.display = 'block'); icon.addEventListener("mouseout", () => icon.querySelectorAll('[class^="icon-tooltip"]')[0].style.display = 'none'); }) } _getNextPost() { return new Promise((resolve) => { const nextHref = qs('.tensorsite-detail__next-url'); if (this.detailBody && nextHref) { let request = new XMLHttpRequest(); request.open('GET', nextHref.getAttribute('href'), true); request.onload = function() { if (this.status >= 200 && this.status < 400) { // Success! Should be an HTML response // Save html in variable so you're able to query select const parser = new DOMParser(); const html = parser.parseFromString(this.response, "text/html"); const nextTitle = html.querySelector('.tensorsite-detail__title'); const nextDesc = html.querySelector('.tensorsite-detail__description'); const nextTags = html.querySelector('.tensorsite-detail__tags'); const nextHref = qs('.tensorsite-detail__next-url').getAttribute('href'); const nextImgUrl = html.querySelector('.tensorsite-detail__main-image'); const nextTitleEl = document.querySelector('.tensorsite-content__title.next'); const nextDescEl = document.querySelector('.tensorsite-content__description.next'); let nextTagsEl = document.querySelector('.tensorsite-content__subtitle.next'); const nextHrefEl = document.querySelector('.tensorsite-card__href.next'); const nextImgEl = document.querySelector('.tensorsite-content__image-wrapper'); const nextContainer = qs('.tensorsite-next'); const footer = qs('.tensorsite-footer'); if (nextTitleEl && nextTitle) { nextTitleEl.innerHTML = nextTitle.innerHTML; } if (nextDescEl && nextDesc) { nextDescEl.innerHTML = nextDesc.innerHTML; } if (nextTagsEl && nextTags) { nextTagsEl.innerHTML = nextTags.innerHTML; } if (nextHref && nextHrefEl) { nextHrefEl.setAttribute('href', nextHref); } if (nextImgEl && nextImgUrl) { // If Blogger can't find a firstImageUrl, it returns a // message informing us of that, so this checks // if the string is a URL if(!/http/.test(nextImgUrl.innerHTML)){ nextImgEl.classList.add('hidden'); } else { nextImgEl.querySelector('img').src = nextImgUrl.innerHTML; } } if (nextHref) { nextContainer.classList.add('active'); footer.classList.add('grey'); } resolve(); } else { // We reached our target server, but it returned an error console.error('Error: Could not get the next title'); } }; request.send(); } }) } get isMenuOpen() { return this.sideMenu.classList.contains('is-open'); } _handleSearchKeypress(e) { if (e.which == 13) { this._searchGoogle(); } } _searchGoogle(e) { e.preventDefault(); const {value} = e.target.querySelector('.search-input'); window.location.href = 'https://www.google.com/search?q=site%3A' + window.location.hostname + '%20' + value; } _toggleMobileMenu() { this.body.classList.toggle('no-scroll'); this.overlay.classList.toggle('show'); this.sideMenu.classList.toggle('is-open'); if (this.isMenuOpen) { this.overlay.addEventListener('click', this._closeMenu); } else { this.overlay.removeEventListener('click', this._closeMenu); } } _closeMenu(e) { if (this.isMenuOpen) { this._toggleMobileMenu(); } } _onResize() { if (this._getScreen().width > 839 && this.isMenuOpen) { this._closeMenu(); } } _getScreen() { return { scrollY: window.scrollY, width: window.innerWidth, height: window.innerHeight, } }; _removeDividerAboveImage() { if (this.detailBody && this.detailBody.firstElementChild && this.detailBody.firstElementChild.querySelector('img')) { const firstDivider = qs('.divider'); firstDivider.style.display = 'none'; } } _isCurrentPathAllPosts(){ const {pathname, search} = window.location; return pathname === '/' || (pathname === '/search' && !/label/.test(search)) } _setAllTagActive() { if(this._isCurrentPathAllPosts()){ const allTag = qs('.header__nav-item.all'); allTag.parentElement.classList.add('active'); } } // Shows featured post only if not on the home page _showFeaturedPost() { if (window.location.pathname === '/' && this.featuredCard) { // If any posts in the list have the same href, hide them. // Using a boolean in order to skip the first match // since that is the actual featured card. let skippedFirst = false; this.cards.forEach(card => { const link = card .querySelector('.tensorsite-card__href') .getAttribute('href'); if (link === this.featuredPostHref) { if (skippedFirst) { card.classList.add('hidden'); } skippedFirst = true; } }) } } _makeImagesZoomable(){ this.images.forEach(image => { image.setAttribute('data-action', 'zoom'); if(/a/i.test(image.parentNode.tagName)){ image.parentNode.replaceWith(image) } }) } // Adds max-results query param if URL contains a label filter but // doesn't contain max-results _redirectWithMaxResults(){ const {search} = window.location; const isLabelMatch = /(tensorflow|tfx|community)/gi.test(search) if(isLabelMatch && !/max-results/.test(search)){ window.location.href = window.location.href + '&max-results=20' } } _removeCardLineBreaks(){ const descriptions = this.cardDescriptions || [this.hiddenDescription]; if(descriptions){ descriptions.forEach(node=> { let stringArray = node.innerText.split(/(\r\n|\n|\r)/g); for(let i = stringArray.length; i > 0; i--){ if(/(\r\n|\n|\r)/.test(stringArray[i]) || stringArray[i] === ''){ let j = i - 1; while(j > 0 && (/(\r\n|\n|\r)/.test(stringArray[j]) || stringArray[j] === '')){ stringArray.splice(j, 1); j-- } } } return node.innerText = stringArray.join('') }) } } } window.addEventListener('DOMContentLoaded', (event) => { new App(); }); //]]> </script> <script type='text/javascript'> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-142343919-1', 'auto', 'blogger'); ga('blogger.send', 'pageview'); </script> <link href='https://www.blogger.com/dyn-css/authorization.css?targetBlogID=7864883956188652345&zx=fec37b9b-5beb-4355-a145-f93f4e783e5d' media='none' onload='if(media!='all')media='all'' rel='stylesheet'/><noscript><link href='https://www.blogger.com/dyn-css/authorization.css?targetBlogID=7864883956188652345&zx=fec37b9b-5beb-4355-a145-f93f4e783e5d' rel='stylesheet'/></noscript> <meta name='google-adsense-platform-account' content='ca-host-pub-1556223355139109'/> <meta name='google-adsense-platform-domain' content='blogspot.com'/> </head> <body> <div class='header__overlay'></div> <div class='section' id='nav'><div class='widget HTML' data-version='1' id='HTML1'> <header class='header'> <div aria-hidden='true' data-href='https://blog.tensorflow.org/' hidden='true' id='home-href'></div> <div class='top-row'> <div class='top-row__left'> <button aria-label='Toggle menu' class='header__hamburger' type='button'> <i class='material-icons'>menu</i> </button> <a class='tensorsite-blog-logo' href='https://blog.tensorflow.org/'> <img alt='TensorFlow Blog Logo' class='tensorsite-blog-logo__image' src='https://www.gstatic.com/tf_blog/images/tfblog_logo.svg'/> </a> </div> <div class='top-row__right'> <div class='header__search-container'> <form action='' class='searchbox'> <input aria-label='Search box' class='search-input' name='q' onblur='if (this.value=="") {this.value="Search the Blog";}' onfocus='if (this.value=="Search the Blog") {this.value=""}' placeholder='Search the Blog' type='text' value=''/> <i class='material-icons'>search</i> <input style='visibility:hidden;position:absolute' type='submit'/> </form> </div> <a class='header__cta' href='https://www.tensorflow.org/'> <svg class='cta-icon' height='12' viewBox='0 0 18 18' width='12' xmlns='http://www.w3.org/2000/svg'> <g fill='none' fill-rule='evenodd' transform='translate(-3 -3)'> <rect height='24' width='24'></rect> <path d='M20.55,10.95 L13.05,3.45 C12.45,2.85 11.55,2.85 10.95,3.45 C10.35,4.05 10.35,4.95 10.95,5.55 L15.9,10.5 L4.5,10.5 C3.6,10.5 3,11.1 3,12 C3,12.9 3.6,13.5 4.5,13.5 L15.9,13.5 L10.95,18.45 C10.35,19.05 10.35,19.95 10.95,20.55 C10.95,20.55 10.95,20.55 10.95,20.55 C11.55,21.15 12.45,21.15 13.05,20.55 C13.05,20.55 13.05,20.55 13.05,20.55 L20.55,13.05 C21.15,12.45 21.15,11.55 20.55,10.95 C20.55,10.95 20.55,10.95 20.55,10.95 Z' fill='#000'></path> </g> </svg> Return to TensorFlow Home </a> </div> </div> <div class='nav-row'> <div class='nav-items'> <tab> <a class='header__nav-item all' dir='ltr' href='https://blog.tensorflow.org/'> All </a> </tab> <tab> <a class='header__nav-item' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow+Core&max-results=20'> TensorFlow Core </a> </tab> <tab> <a class='header__nav-item' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow.js&max-results=20'> TensorFlow.js </a> </tab> <tab> <a class='header__nav-item' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow+Lite&max-results=20'> TensorFlow Lite </a> </tab> <tab> <a class='header__nav-item' dir='ltr' href='https://blog.tensorflow.org/search?label=TFX&max-results=20'> TFX </a> </tab> <tab> <a class='header__nav-item' dir='ltr' href='https://blog.tensorflow.org/search?label=Community&max-results=20'> Community </a> </tab> </div> <section class='social-icons'> <div class='social-icons__container-header'> <div class='social-icons__links'> <a class='icon-link' href='https://discuss.tensorflow.org' rel='noopener noreferrer' target='_blank'> <img alt='TensorFlow Forum' src='https://www.gstatic.com/tf_blog/images/ic_forum_2.svg'/> <div class='icon-tooltip'>TensorFlow Forum</div> </a> <a class='icon-link' href='https://www.youtube.com/tensorflow' rel='noopener noreferrer' target='_blank'> <img alt='TensorFlow YouTube' src='https://www.gstatic.com/tf_blog/images/ic_youtube.svg'/> <div class='icon-tooltip'>TensorFlow YouTube</div> </a> <a class='icon-link' href='https://twitter.com/TensorFlow' rel='noopener noreferrer' target='_blank'> <img alt='TensorFlow Twitter' src='https://www.gstatic.com/tf_blog/images/ic_twitter.svg'/> <div class='icon-tooltip'>TensorFlow Twitter</div> </a> <a class='icon-link' href='https://github.com/tensorflow' rel='noopener noreferrer' target='_blank'> <img alt='TensorFlow GitHub' src='https://www.gstatic.com/tf_blog/images/ic_github.svg'/> <div class='icon-tooltip-github'>TensorFlow GitHub</div> </a> </div> </div> </section> </div> </header> <div class='header__side-menu'> <div class='header__side-menu__logo-container'> <a class='tensorsite-blog-logo' href='https://blog.tensorflow.org/'> <img alt='TensorFlow Blog Logo' class='tensorsite-blog-logo__image' src='https://www.gstatic.com/tf_blog/images/tfblog_logo.svg'/> </a> </div> <div class='header__side-menu__content'> <div class='header__side-menu__items'> <div class='header__search-container mobile'> <form action='' class='searchbox'> <input aria-label='Search box' class='search-input' name='q' onblur='if (this.value=="") {this.value="Search the Blog";}' onfocus='if (this.value=="Search the Blog") {this.value=""}' placeholder='Search the Blog' type='text' value=''/> <i class='material-icons'>search</i> <input style='visibility:hidden;position:absolute' type='submit'/> </form> </div> <div class='header__side-menu__title'>Tags</div> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/'> All </a> </li> </tab> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow+Core&max-results=20'> TensorFlow Core </a> </li> </tab> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow.js&max-results=20'> TensorFlow.js </a> </li> </tab> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/search?label=TensorFlow+Lite&max-results=20'> TensorFlow Lite </a> </li> </tab> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/search?label=TFX&max-results=20'> TFX </a> </li> </tab> <tab> <li> <a class='header__nav-item mobile' dir='ltr' href='https://blog.tensorflow.org/search?label=Community&max-results=20'> Community </a> </li> </tab> </div> <div class='spacer'></div> <div class='header__side-menu__bottom'> <a class='header__cta mobile' href='https://www.tensorflow.org/'> Return to TensorFlow Home </a> </div> </div> </div> </div></div> <div class='content-wrap'> <div class='section' id='blog'><div class='widget FeaturedPost' data-version='1' id='FeaturedPost1'> </div><div class='widget Blog' data-version='1' id='Blog1'> <div class='tensorsite-container--narrow'> <div class='tensorsite-detail'> <a aria-hidden='true' class='tensorsite-detail__next-url' hidden='true' href='https://blog.tensorflow.org/2022/09/fast-reduce-and-mean-in-tensorflow-lite.html'></a> <div aria-hidden='true' class='tensorsite-detail__current-url' hidden='true'>https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html</div> <div aria-hidden='true' class='tensorsite-detail__tags' hidden='true'> <span>Flax</span> <b class='label-divider-dot'>·</b> <span>JAX</span> </div> <div aria-hidden='true' class='tensorsite-detail__main-image' hidden='true'> https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png </div> <p aria-hidden='true' class='tensorsite-detail__description' hidden='true'> <span class='tensorsite-content__info'> October 03, 2022 — </span> <em>Posted by Wei Wei, Developer Advocate</em>In our previous blog post <a href="https://blog.tensorflow.org/2021/10/building-board-game-app-with-tensorflow.html" target="_blank">Building a board game app with TensorFlow: a new TensorFlow Lite reference app</a>, we showed you how to use TensorFlow and TensorFlow Agents to train a reinforcement learning (RL) agent to play a simple board game ‘Plane Strike’. We also converted the trained model to TensorFlow Lite and then deployed it into a fully-functional Android a… </p> <div class='tensorsite-content__subtitle'> <a href='https://blog.tensorflow.org/search?label=Flax&max-results=20'> <span>Flax</span> </a> <b class='label-divider-dot'>·</b> <a href='https://blog.tensorflow.org/search?label=JAX&max-results=20'> <span>JAX</span> </a> <b class='label-divider-dot'>·</b> <img alt='Google Article' class='community-icon' src='https://www.gstatic.com/tf_blog/images/ic_google.svg'/> </div> <div class='tensorsite-detail__title'> Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite </div> <div class='tensorsite-detail__contact'> <div class='tensorsite-detail__info'> <span class='tensorsite-detail__timestamp'>October 03, 2022</span> </div> <a class='icon-link' href='https://twitter.com/intent/tweet?text=%22Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite %22 from the TensorFlow Blog%0A%0Ahttps://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html' rel='noopener noreferrer' target='_blank' title='Share this post on Twitter'> <svg alt='Twitter Social Icon' class='twitter-icon social-icon' height='19' viewBox='0 0 23 19' width='23' xmlns='http://www.w3.org/2000/svg'> <g fill='none' fill-rule='evenodd' transform='translate(-7 -9)'> <rect height='36' width='36'></rect> <path d='M14.076,27.2827953 C22.566,27.2827953 27.21,20.2477953 27.21,14.1477953 C27.21,13.9477953 27.21,13.7487953 27.197,13.5507953 C28.1,12.8977953 28.88,12.0887953 29.5,11.1617953 C28.657,11.5347953 27.764,11.7797953 26.848,11.8877953 C27.812,11.3107953 28.533,10.4037953 28.878,9.33479527 C27.972,9.87179527 26.98,10.2507953 25.947,10.4547953 C24.198,8.59579527 21.274,8.50679527 19.415,10.2547953 C18.217,11.3817953 17.708,13.0617953 18.08,14.6647953 C14.368,14.4787953 10.91,12.7257953 8.566,9.84279527 C7.341,11.9507953 7.967,14.6497953 9.995,16.0047953 C9.261,15.9827953 8.542,15.7837953 7.9,15.4267953 L7.9,15.4847953 C7.9,17.6827953 9.449,19.5747953 11.603,20.0107953 C10.924,20.1957953 10.211,20.2227953 9.519,20.0897953 C10.124,21.9707953 11.856,23.2587953 13.832,23.2957953 C12.197,24.5797953 10.178,25.2777953 8.098,25.2747953 C7.731,25.2747953 7.364,25.2527953 7,25.2087953 C9.111,26.5627953 11.567,27.2817953 14.076,27.2787953' fill='#545454'></path> </g> </svg> </a> </div> <div class='divider divider--article-top'></div> <div class='tensorsite-detail__body'> <meta content="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png" name="twitter:image"></meta> <img src="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png" style="display: none;" /> <p><em>Posted by Wei Wei, Developer Advocate</em></p><p> <a href="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEi5XPGcGH5styeLZUJPU4pF0VhSor6M2lT_nGVjsid6Nr-tOuObab9Pqwsn83FiOO8XUS2b-g8ugd2BnwAqTFjKxxUiDZaIabDMxVeO4Bvl4Y7EduMOCQ4IqU3JgKaa3fJUWUVfmkWuGZ8dsH2EtDqe2gj8NMw4GfhkZe1bYQRdTTY3sauYMLn1xaqn/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-header.png"><img border="0" data-original-height="800" data-original-width="1058" height="198" src="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEi5XPGcGH5styeLZUJPU4pF0VhSor6M2lT_nGVjsid6Nr-tOuObab9Pqwsn83FiOO8XUS2b-g8ugd2BnwAqTFjKxxUiDZaIabDMxVeO4Bvl4Y7EduMOCQ4IqU3JgKaa3fJUWUVfmkWuGZ8dsH2EtDqe2gj8NMw4GfhkZe1bYQRdTTY3sauYMLn1xaqn/w666-h198/Tensorflow-building-reinforcement-learning-agent-with-JAX-header.png" width="666" /></a> </p><a name='more'></a><p></p> <div><p>In our previous blog post <a href="https://blog.tensorflow.org/2021/10/building-board-game-app-with-tensorflow.html" target="_blank">Building a board game app with TensorFlow: a new TensorFlow Lite reference app</a>, we showed you how to use TensorFlow and TensorFlow Agents to train a reinforcement learning (RL) agent to play a simple board game ‘Plane Strike’. We also converted the trained model to TensorFlow Lite and then deployed it into a fully-functional Android app. In this blog, we will demonstrate a new path: train the same RL agent with <a href="https://flax.readthedocs.io/" target="_blank">Flax</a>/<a href="https://jax.readthedocs.io/" target="_blank">JAX</a> and deploy it into the same Android app we have built before. The complete code has been open sourced in the <a href="https://github.com/tensorflow/examples/blob/master/lite/examples/reinforcement_learning/ml/tf_and_jax/training_jax.py" target="_blank">tensorflow/examples</a> repository for your reference.<br /><br />To refresh your memory, our RL-based agent needs to predict a strike position based on the human player’s board position so that it can finish the game before the human player does. For more detailed game rules, please refer to our previous <a href="https://blog.tensorflow.org/2021/10/building-board-game-app-with-tensorflow.html" target="_blank">blog</a>. </p><table align="center" cellpadding="0" cellspacing="0" class="tr-caption-container" style="margin-left: auto; margin-right: auto; text-align: center;"><tbody><tr><td style="text-align: center;"><center><img alt="Demo game play in ‘Plane Strike’" border="0" data-original-height="1504" data-original-width="720" src="https://blogger.googleusercontent.com/img/a/AVvXsEj7K9sQBJ8GVQcgPNGxsMGJmgCkoPaOxvbTJttSqn0kRxpRDhVSW7gPzr93vcDWOhZDY73YKOSG0_ERiJHwcw9T08EuwUvUXUAt7bFC8giFOu-Shl6FWnWmILajHGZ2K41XSzCtKNoiEpMCc2WnxGavNPac2Ua8T2iGemtud1NQ_pVfRutYVlzv6HZt" style="width: 50%;" td="" /></center></td></tr> <tr><td class="tr-caption" style="text-align: center;">Demo game play in ‘Plane Strike’</td></tr></tbody></table></div><div><h3 style="text-align: left;">Background: JAX and TensorFlow</h3><a href="https://github.com/google/jax" target="_blank">JAX</a> is a NumPy-like library developed by Google Research for high performance computing. It uses <a href="https://www.tensorflow.org/xla" target="_blank">XLA</a> to compile programs optimized for GPUs and <a href="https://cloud.google.com/tpu" target="_blank">TPUs</a>. <a href="https://github.com/google/flax" target="_blank">Flax</a> is a popular neural network library built on top of JAX. Researchers have been using JAX/Flax to train very large models with billions of parameters (such as <a href="https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html" target="_blank">PaLM</a> for language understanding and generation, or <a href="https://imagen.research.google/" target="_blank">Imagen</a> for image generation), making full use of modern hardware. If you're new to JAX and Flax, start with <a href="https://jax.readthedocs.io/en/latest/jax-101/index.html" target="_blank">this JAX 101 tutorial</a> and <a href="https://flax.readthedocs.io/en/latest/getting_started.html" target="_blank">this Flax Getting Started example</a>.<br /><br />TensorFlow started as a library for ML towards the end of 2015 and has since become a rich ecosystem that includes tools for productionizing ML pipelines (<a href="https://www.tensorflow.org/tfx" target="_blank">TFX</a>), data visualization (<a href="https://tensorboard.dev/" target="_blank">TensorBoard</a>), deploying ML models to edge devices (<a href="https://www.tensorflow.org/lite" target="_blank">TensorFlow Lite</a>), and devices running on a web browser or any device capable of executing JavaScript (<a href="https://www.tensorflow.org/js" target="_blank">TensorFlow.js</a>). Models developed in JAX or Flax can tap into this rich ecosystem by first converting such a model to the TensorFlow <a href="https://www.tensorflow.org/guide/saved_model" target="_blank">SavedModel</a> format, and then using the same tooling as if they had been developed in TensorFlow natively.</div><div><br /></div><div>If you already have a JAX-trained model and want to deploy it today, we have put together a list of resources for you:<br /><ul style="text-align: left;"><li><a href="https://youtu.be/I4dx7OI9FJQ?t=36" target="_blank">Serving JAX models with TensorFlow Serving</a> video shows you how to deploy JAX models with TensorFlow Serving</li></ul><ul style="text-align: left;"><li><a href="https://blog.tensorflow.org/2022/08/jax-on-web-with-tensorflowjs.html" target="_blank">JAX on the Web with TensorFlow.js</a> blog walks through how to convert JAX models to TFJS and run them within web apps</li></ul><ul style="text-align: left;"><li>This blog post demos how to convert a Flax/JAX model to TFLite and run it in a native Android app</li></ul>Overall, no matter what your deployment target is (server, web or mobile), we got you covered.<br />Implementing the game agent with Flax/JAX<br /><br />Coming back to our board game, to implement our RL agent, we will leverage the same <a href="https://github.com/tensorflow/examples/tree/master/lite/examples/reinforcement_learning/ml/tf_and_jax/gym_planestrike/gym_planestrike/envs" target="_blank">gym</a> environment as before. We will train the same policy gradient model using Flax/JAX this time. Recall that mathematically the policy gradient is defined as:</div><div>  <br /><div class="separator" style="clear: both; text-align: center;"><a href="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEjaOdJkvZr0cr1-e-5ZrmcejmyUcrtp6BH82QaEQyehH02rUXCdIpl6yj5FEh8brzWV2FlVqrdDBVYwejA25B5cmk3spWoK0cSLjBCdq8QPL9G45dVfl_ZlPvyVGoXpvmc3x4X44E4Q2F-DrwinPw1AMWFckqrYjvT3s-a4k3F6Alv-qdMjCDfKyHMS/s512/unnamed%20(12).png" style="margin-left: 1em; margin-right: 1em;"><img border="0" data-original-height="151" data-original-width="512" src="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEjaOdJkvZr0cr1-e-5ZrmcejmyUcrtp6BH82QaEQyehH02rUXCdIpl6yj5FEh8brzWV2FlVqrdDBVYwejA25B5cmk3spWoK0cSLjBCdq8QPL9G45dVfl_ZlPvyVGoXpvmc3x4X44E4Q2F-DrwinPw1AMWFckqrYjvT3s-a4k3F6Alv-qdMjCDfKyHMS/s16000/unnamed%20(12).png" /></a></div></div><p></p><p></p> <p> where: </p> <ul> <li><em>T</em>: the number of timesteps per episode, which can vary per episode </li><li><em>s<sub>t</sub></em>: the state at timestep <em>t</em> </li><li><em>a<sub>t</sub></em>: chosen action at timestep <em>t </em>given state <em>s</em> </li><li><em>蟺<sub>胃</sub></em>: the policy parameterized by <em>胃</em> </li><li><em>R(*)</em>: the reward gathered, given the policy </li> </ul> <p> We define a 3-layer MLP as our policy network, which predicts the agent’s next strike position. </p><p><span id="docs-internal-guid-5153e9fe-7fff-c196-ded1-82135867f5f2"></span></p><div align="left" dir="ltr" style="margin-left: 0pt;"><table style="border-collapse: collapse; border: none;"><colgroup><col width="719"></col></colgroup><tbody><tr style="height: 0pt;"><td style="background-color: #efecf4; overflow-wrap: break-word; overflow: hidden; padding: 5pt; vertical-align: top;"><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: blue; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">class</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #267f99; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">PolicyGradient</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #267f99; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">nn</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">.</span><span style="background-color: transparent; color: #267f99; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">Module</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">"""Neural network to predict the next strike position."""</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">@nn.compact</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: blue; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">__call__</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">self</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">x</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    dtype = jnp.float32</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">    x = x.reshape((x.shape[</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">0</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">], -</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">1</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">))</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    x = nn.Dense(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">        </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">features</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">2</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"> * common.BOARD_SIZE**</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">2</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">name</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">'hidden1'</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">dtype</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=dtype)(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">           x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    x = nn.relu(x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">    x = nn.Dense(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">features</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=common.BOARD_SIZE**</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">2</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">name</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">'hidden2'</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">dtype</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=dtype)(x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    x = nn.relu(x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">    x = nn.Dense(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">features</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=common.BOARD_SIZE**</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">2</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">name</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">'logits'</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">dtype</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">=dtype)(x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    policy_probabilities = nn.softmax(x)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">    </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: black; font-size: 9pt; font-style: normal; font-variant: normal; font-weight: 400; text-decoration: none; vertical-align: baseline; white-space: pre-wrap;"> policy_probabilities</span></span></p></td></tr></tbody></table><br /></div>In our main training loop, in each iteration we use the neural network to play a round of the game, gather the trajectory information (game board positions, actions taken and rewards), discount the rewards, and then train the model with the trajectories. <br /> <div><br /></div><div><span id="docs-internal-guid-986ac9b3-7fff-6d43-2dd4-c3dbbf30f52e"><div align="left" dir="ltr" style="margin-left: 0pt;"><table style="border-collapse: collapse; border: none;"><colgroup><col width="719"></col></colgroup><tbody><tr style="height: 0pt;"><td style="background-color: #efecf4; overflow-wrap: break-word; overflow: hidden; padding: 5pt; vertical-align: top;"><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">for</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> i </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">in</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> tqdm(</span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">range</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(iterations)):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">   predict_fn = functools.partial(run_inference, params)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">   board_log, action_log, result_log = common.play_game(predict_fn)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">   rewards = common.compute_rewards(result_log)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">   optimizer, params, opt_state = train_step(optimizer, params, opt_state,</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">                                             board_log, action_log, rewards)</span></span></p></td></tr></tbody></table><br /></div>In the train_step() method, we first compute the loss using the trajectories. Then we use jax.grad() to compute the gradients. Lastly we use <a href="https://github.com/deepmind/optax" target="_blank">Optax</a>, a gradient processing and optimization library for JAX, to update the model parameters. </span></div><div><span><br /></span></div><div><span id="docs-internal-guid-aacd0307-7fff-1bc1-a623-b469cd177dae"><div align="left" dir="ltr" style="margin-left: 0pt;"><table style="border-collapse: collapse; border: none;"><colgroup><col width="719"></col></colgroup><tbody><tr style="height: 0pt;"><td style="background-color: #efecf4; overflow-wrap: break-word; overflow: hidden; padding: 5pt; vertical-align: top;"><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">compute_loss</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">logits</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">labels</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">rewards</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  one_hot_labels = jax.nn.one_hot(labels, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">num_classes</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=common.BOARD_SIZE**</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">2</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">  loss = -jnp.mean(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">      jnp.sum(one_hot_labels * jnp.log(logits), </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">axis</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=-</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">1</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">) * jnp.asarray(rewards))</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> loss</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">train_step</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">model_optimizer</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">params</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">opt_state</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">game_board_log</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">,</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">              </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">predicted_action_log</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">action_result_log</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">"""Run one training step."""</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">loss_fn</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">model_params</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    logits = run_inference(model_params, game_board_log)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">    loss = compute_loss(logits, predicted_action_log, action_result_log)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">    </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> loss</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">compute_grads</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">params</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">    </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> jax.grad(loss_fn)(params)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">  grads = compute_grads(params)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">  updates, opt_state = model_optimizer.update(grads, opt_state)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">  params = optax.apply_updates(params, updates)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> model_optimizer, params, opt_state</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">@jax.jit</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">def</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">run_inference</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(</span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">model_params</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">board</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">):</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  logits = PolicyGradient().apply({</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">'params'</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">: model_params}, board)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.62857; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">return</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> logits</span></span></p></td></tr></tbody></table></div><br />That’s it for the training loop. We can visualize the training progress in TensorBoard as below; here we use the proxy metric ‘game_length’ (the number of steps to finish the game) to track the progress. The intuition is that when the agent becomes smarter, it can finish the game in fewer steps.</span></div><div><span><br /></span></div><div><div class="separator" style="clear: both; text-align: center;"><a href="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEi1O2B--pbmUaBlRBbCR_5ug9zeIj9vD6W7iwokGbvTKHo1l3hIgbYcZOcugSnskgXtU4tIvfuEsgncv-L5I050wHOaytu6o5xVLVQy4b6ilO9IoqzfesjZWY1HD5bCs_785xbjjFUOkyTuwpBewnKo1tHGxo7-PwmBEBFkjR40q3QrfV23d9sM5Pad/s1043/image2.png" style="margin-left: 1em; margin-right: 1em;"><img border="0" data-original-height="534" data-original-width="1043" src="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEi1O2B--pbmUaBlRBbCR_5ug9zeIj9vD6W7iwokGbvTKHo1l3hIgbYcZOcugSnskgXtU4tIvfuEsgncv-L5I050wHOaytu6o5xVLVQy4b6ilO9IoqzfesjZWY1HD5bCs_785xbjjFUOkyTuwpBewnKo1tHGxo7-PwmBEBFkjR40q3QrfV23d9sM5Pad/s16000/image2.png" /></a></div><h3 style="text-align: left;">Converting the Flax/JAX model to TensorFlow Lite and integrating with the Android app</h3>After the model is trained, we use the <a href="https://github.com/google/jax/tree/main/jax/experimental/jax2tf" target="_blank">jax2tf</a>, a TensorFlow-JAX interoperation tool, to convert the JAX model into a TensorFlow concrete function. And the final step is to call TensorFlow Lite converter to convert the concrete function into a TFLite model. </div><div><br /></div><div><span id="docs-internal-guid-f6dd41f2-7fff-1eee-84c6-e3fc47a283e8"><div align="left" dir="ltr" style="margin-left: 0pt;"><table style="border-collapse: collapse; border: none;"><colgroup><col width="719"></col></colgroup><tbody><tr style="height: 0pt;"><td style="background-color: #efecf4; overflow-wrap: break-word; overflow: hidden; padding: 5pt; vertical-align: top;"><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: green; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"># Convert to tflite model</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;"> model = PolicyGradient()</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> jax_predict_fn = </span><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">lambda</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">input</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">: model.apply({</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">'params'</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">: params}, </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">input</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;"> tf_predict = tf.function(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">     jax2tf.convert(jax_predict_fn, </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">enable_xla</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">False</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">),</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">     </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">input_signature</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=[</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">         tf.TensorSpec(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">             </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">shape</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=[</span><span style="background-color: transparent; color: #098658; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">1</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">, common.BOARD_SIZE, common.BOARD_SIZE],</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">             </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">dtype</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=tf.float32,</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">             </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">name</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">'input'</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">     ],</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">     </span><span style="background-color: transparent; color: #001080; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">autograph</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">=</span><span style="background-color: transparent; color: blue; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">False</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">,</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;"> )</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;"> converter = tf.lite.TFLiteConverter.from_concrete_functions(</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">     [tf_predict.get_concrete_function()], tf_predict)</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;"> tflite_model = converter.convert()</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"> </span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: green; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"># Save the model</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">with</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> </span><span style="background-color: transparent; color: #795e26; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">open</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">(os.path.join(modeldir, </span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">'planestrike.tflite'</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">), </span><span style="background-color: transparent; color: #a31515; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">'wb'</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">) </span><span style="background-color: transparent; color: #af00db; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">as</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> f:</span></span></p><p dir="ltr" style="background-color: white; line-height: 1.8; margin-bottom: 0pt; margin-top: 0pt;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><span style="font-family: courier;">   f.write(tflite_model)</span></span></p></td></tr></tbody></table><br /></div>The JAX-converted TFLite model behaves exactly like any TensorFlow-trained TFLite model. You can visualize it with Netron:</span></div> <table align="center" cellpadding="0" cellspacing="0" class="tr-caption-container" style="margin-left: auto; margin-right: auto; text-align: center;"><tbody><tr><td style="text-align: center;"><center><img alt="Visualizing TFLite model converted from Flax/JAX using Netron" border="0" data-original-height="1504" data-original-width="720" src="https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEg0gjG2lKGAENN2Y9d0WH3WyGY6Grwx09hel-yRK0bkfD3KjGnsmLEYIFpfexazqhSvmuZ7HIPOmFKqF1cw6Qq4VWzWK4C6v4pYg0aMBcHjbwVnlfDW5TnYPmJRm48jAwtlLcCFGjy5lT6rJ0Pi4cIaAfUxFkNiQJ-NQSm347xO9nDpKxV4L_Y22-oB/s1600/image1.png" style="width: 50%;" td="" /></center></td></tr> <tr><td class="tr-caption" style="text-align: center;">Visualizing TFLite model converted from Flax/JAX using Netron</td></tr></tbody></table> <div>We can use exactly the same Java code as before to invoke the model and get the prediction.</div><div><span id="docs-internal-guid-a7854efe-7fff-d283-d5ca-b5d3c2794d5e"><p dir="ltr" style="background-color: white; line-height: 1.38; margin-bottom: 9pt; margin-top: 0pt;"></p><div align="left" dir="ltr" style="margin-left: 0pt;"><table style="border-collapse: collapse; border: none;"><colgroup><col width="719"></col></colgroup><tbody><tr style="height: 0pt;"><td style="background-color: #efecf4; overflow-wrap: break-word; overflow: hidden; padding: 5pt; vertical-align: top;"><p dir="ltr" style="line-height: 1.38; margin-bottom: 0pt; margin-top: 0pt;"><span style="font-family: courier;"><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">convertBoardStateToByteBuffer(board);</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">tflite.run(boardData, outputProbArrays);</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">float[] probArray = outputProbArrays[</span><span style="background-color: transparent; color: #aa573c; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">0</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">];</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">int agentStrikePosition = </span><span style="background-color: transparent; color: #aa573c; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">-1</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">;</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">float maxProb = </span><span style="background-color: transparent; color: #aa573c; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">0</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">;</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #955ae7; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">for</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> (int i = </span><span style="background-color: transparent; color: #aa573c; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">0</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">; i < probArray.length; i++) {</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  int x = i / Constants.BOARD_SIZE;</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  int y = i % Constants.BOARD_SIZE;</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  </span><span style="background-color: transparent; color: #955ae7; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">if</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"> (board[x][y] == BoardCellStatus.UNTRIED && probArray[i] > maxProb) {</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">    agentStrikePosition = i;</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">    maxProb = probArray[i];</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">  }</span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;"><br /></span><span style="background-color: transparent; color: #585260; font-size: 9pt; font-variant-east-asian: normal; font-variant-numeric: normal; vertical-align: baseline; white-space: pre-wrap;">}</span></span></p></td></tr></tbody></table></div></span><h3 style="text-align: left;"><span>Conclusion</span></h3><span>In summary, this article walks you through how to train a simple reinforcement learning model with Flax/JAX, leverage <a href="https://github.com/google/jax/tree/main/jax/experimental/jax2tf" target="_blank">jax2tf</a> to convert it to TensorFlow Lite, and integrate the converted model into an Android app.<br /><br />Now you have learned how to build neural network models with Flax/JAX, and tap into the powerful TensorFlow ecosystem to deploy your models pretty much anywhere you want. We can’t wait to see the fantastic apps you build with both JAX and TensorFlow!</span></div> </div> </div> <div class='tensorsite-detail-footer'> <div class='article-divider'> <img alt='Diamond Article Divider' src='https://www.gstatic.com/tf_blog/images/ic_article_end.svg'/> </div> <a class='tensorsite-chip' href='https://blog.tensorflow.org/search?label=Flax&max-results=20'> Flax </a> <a class='tensorsite-chip' href='https://blog.tensorflow.org/search?label=JAX&max-results=20'> JAX </a> <a class='tensorsite-chip' href='https://blog.tensorflow.org/search?label=TensorFlow+Lite&max-results=20'> TensorFlow Lite </a> </div> </div> <div class='tensorsite-next'> <div class='tensorsite-container--large'> <div class='tensorsite-next__title'>Next post</div> <div class='tensorsite-card'> <a aria-label='Next Card' class='tensorsite-card__href next' href='https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html'></a> <div class='tensorsite-content__image-wrapper'> <img alt='Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ' class='tensorsite-content__image' src='https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png'/> </div> <div class='tensorsite-content'> <div class='tensorsite-content__subtitle next'> <span>Flax</span> <b class='label-divider-dot'>·</b> <span>JAX</span> <b class='label-divider-dot'>·</b> <img alt='Google Article' class='community-icon' src='https://www.gstatic.com/tf_blog/images/ic_google.svg'/> </div> <div class='tensorsite-content__title next'> Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite </div> <p class='tensorsite-content__description next'> <span class='tensorsite-content__info'> October 03, 2022 </span> — <span> <em>Posted by Wei Wei, Developer Advocate</em>In our previous blog post <a href="https://blog.tensorflow.org/2021/10/building-board-game-app-with-tensorflow.html" target="_blank">Building a board game app with TensorFlow: a new TensorFlow Lite reference app</a>, we showed you how to use TensorFlow and TensorFlow Agents to train a reinforcement learning (RL) agent to play a simple board game ‘Plane Strike’. We also converted the trained model to TensorFlow Lite and then deployed it into a fully-functional Android a… </span> </p> </div> </div> </div> </div>  </div></div> </div>  <div class='tensorsite-full-footer'> <div class='section' id='footer'><div class='widget HTML' data-version='1' id='HTML2'> <section class='tensorsite-footer'> <div class='tensorsite-container tensorsite-footer__container'> <div class='tensorsite-footer__side tensorsite-footer__side--left'></div> <div class='tensorsite-footer__side tensorsite-footer__side--right'></div> <div class='tensorsite-footer__content'> <div class='tensorsite-content'> <div class='tensorsite-content__title tensorsite-content__title--grow'> Build, deploy, and experiment easily with TensorFlow </div> <div class='tensorsite-content__cta-wrapper'> <a class='tensorsite-content__button tensorsite-button tensorsite-button--white' href='https://www.tensorflow.org/'> Get started </a> </div> </div> </div> </div> <div class='tensorsite-footer__lines'></div> </section> <section class='social-icons'> <div class='social-icons__container-footer'> <a alt='TensorFlow Home' class='tensorsite-logo' href='https://www.tensorflow.org/'> <img alt='TensorFlow Logo' class='tensorsite-logo__image' src='https://www.gstatic.com/tf_blog/images/tf_lockup.svg'/> </a> <div class='footer__side--right'> <div class='footer__links'> <a alt='Youtube Social Link' class='footer-link' href='https://www.google.com/' rel='noopener noreferrer' target='_blank'> Google </a> <a alt='Twitter Social Link' class='footer-link' href='https://policies.google.com/privacy' rel='noopener noreferrer' target='_blank'> Privacy </a> <a alt='Github Link' class='footer-link' href='https://policies.google.com/terms' rel='noopener noreferrer' target='_blank'> Terms </a> <a class='footer-link' href='https://blog.tensorflow.org/p/tensorflow-blog-contribution-notice.html'> Contributions notice </a> </div> <div class='social-icons__links'> <a alt='Youtube Social Link' class='icon-link' href='https://www.youtube.com/channel/UC0rqucBdTuFTjJiefW5t-IQ' rel='noopener noreferrer' target='_blank'> <img src='https://www.gstatic.com/tf_blog/images/ic_youtube.svg'/></a> <a alt='Twitter Social Link' class='icon-link' href='https://twitter.com/TensorFlow' rel='noopener noreferrer' target='_blank'> <img src='https://www.gstatic.com/tf_blog/images/ic_twitter.svg'/></a> <a alt='Github Link' class='icon-link' href='https://github.com/tensorflow' rel='noopener noreferrer' target='_blank'> <img src='https://www.gstatic.com/tf_blog/images/ic_github.svg'/></a> </div> </div> </div> </section> </div></div> </div> <script type="text/javascript" src="https://www.blogger.com/static/v1/widgets/984859869-widgets.js"></script> <script type='text/javascript'> window['__wavt'] = 'AOuZoY6c_jaz7MEFHRNtmcUATEoGl7ZXjw:1732374352279';_WidgetManager._Init('//www.blogger.com/rearrange?blogID\x3d7864883956188652345','//blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html','7864883956188652345'); _WidgetManager._SetDataContext([{'name': 'blog', 'data': {'blogId': '7864883956188652345', 'title': 'The TensorFlow Blog', 'url': 'https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html', 'canonicalUrl': 'https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html', 'homepageUrl': 'https://blog.tensorflow.org/', 'searchUrl': 'https://blog.tensorflow.org/search', 'canonicalHomepageUrl': 'https://blog.tensorflow.org/', 'blogspotFaviconUrl': 'https://blog.tensorflow.org/favicon.ico', 'bloggerUrl': 'https://www.blogger.com', 'hasCustomDomain': true, 'httpsEnabled': true, 'enabledCommentProfileImages': true, 'gPlusViewType': 'FILTERED_POSTMOD', 'adultContent': false, 'analyticsAccountNumber': 'UA-142343919-1', 'encoding': 'UTF-8', 'locale': 'en', 'localeUnderscoreDelimited': 'en', 'languageDirection': 'ltr', 'isPrivate': false, 'isMobile': false, 'isMobileRequest': false, 'mobileClass': '', 'isPrivateBlog': false, 'isDynamicViewsAvailable': true, 'feedLinks': '\x3clink rel\x3d\x22alternate\x22 type\x3d\x22application/atom+xml\x22 title\x3d\x22The TensorFlow Blog - Atom\x22 href\x3d\x22https://blog.tensorflow.org/feeds/posts/default\x22 /\x3e\n\x3clink rel\x3d\x22alternate\x22 type\x3d\x22application/rss+xml\x22 title\x3d\x22The TensorFlow Blog - RSS\x22 href\x3d\x22https://blog.tensorflow.org/feeds/posts/default?alt\x3drss\x22 /\x3e\n\x3clink rel\x3d\x22service.post\x22 type\x3d\x22application/atom+xml\x22 title\x3d\x22The TensorFlow Blog - Atom\x22 href\x3d\x22https://www.blogger.com/feeds/7864883956188652345/posts/default\x22 /\x3e\n\n\x3clink rel\x3d\x22alternate\x22 type\x3d\x22application/atom+xml\x22 title\x3d\x22The TensorFlow Blog - Atom\x22 href\x3d\x22https://blog.tensorflow.org/feeds/6024940900969663155/comments/default\x22 /\x3e\n', 'meTag': '', 'adsenseHostId': 'ca-host-pub-1556223355139109', 'adsenseHasAds': true, 'adsenseAutoAds': false, 'boqCommentIframeForm': true, 'loginRedirectParam': '', 'view': '', 'dynamicViewsCommentsSrc': '//www.blogblog.com/dynamicviews/4224c15c4e7c9321/js/comments.js', 'dynamicViewsScriptSrc': '//www.blogblog.com/dynamicviews/d78375fb222d99b3', 'plusOneApiSrc': 'https://apis.google.com/js/platform.js', 'disableGComments': true, 'interstitialAccepted': false, 'sharing': {'platforms': [{'name': 'Get link', 'key': 'link', 'shareMessage': 'Get link', 'target': ''}, {'name': 'Facebook', 'key': 'facebook', 'shareMessage': 'Share to Facebook', 'target': 'facebook'}, {'name': 'BlogThis!', 'key': 'blogThis', 'shareMessage': 'BlogThis!', 'target': 'blog'}, {'name': 'X', 'key': 'twitter', 'shareMessage': 'Share to X', 'target': 'twitter'}, {'name': 'Pinterest', 'key': 'pinterest', 'shareMessage': 'Share to Pinterest', 'target': 'pinterest'}, {'name': 'Email', 'key': 'email', 'shareMessage': 'Email', 'target': 'email'}], 'disableGooglePlus': true, 'googlePlusShareButtonWidth': 0, 'googlePlusBootstrap': '\x3cscript type\x3d\x22text/javascript\x22\x3ewindow.___gcfg \x3d {\x27lang\x27: \x27en\x27};\x3c/script\x3e'}, 'hasCustomJumpLinkMessage': false, 'jumpLinkMessage': 'Read more', 'pageType': 'item', 'postId': '6024940900969663155', 'postImageThumbnailUrl': 'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s72-c/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png', 'postImageUrl': 'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png', 'pageName': 'Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ', 'pageTitle': 'The TensorFlow Blog: Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ', 'metaDescription': 'In this blog post, we will show you how to train a game agent using reinforcement learning using JAX/Flax, convert the model to TensorFlow Lite, and d'}}, {'name': 'features', 'data': {}}, {'name': 'messages', 'data': {'edit': 'Edit', 'linkCopiedToClipboard': 'Link copied to clipboard!', 'ok': 'Ok', 'postLink': 'Post Link'}}, {'name': 'template', 'data': {'name': 'custom', 'localizedName': 'Custom', 'isResponsive': false, 'isAlternateRendering': false, 'isCustom': true}}, {'name': 'view', 'data': {'classic': {'name': 'classic', 'url': '?view\x3dclassic'}, 'flipcard': {'name': 'flipcard', 'url': '?view\x3dflipcard'}, 'magazine': {'name': 'magazine', 'url': '?view\x3dmagazine'}, 'mosaic': {'name': 'mosaic', 'url': '?view\x3dmosaic'}, 'sidebar': {'name': 'sidebar', 'url': '?view\x3dsidebar'}, 'snapshot': {'name': 'snapshot', 'url': '?view\x3dsnapshot'}, 'timeslide': {'name': 'timeslide', 'url': '?view\x3dtimeslide'}, 'isMobile': false, 'title': 'Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite ', 'description': 'In this blog post, we will show you how to train a game agent using reinforcement learning using JAX/Flax, convert the model to TensorFlow Lite, and d', 'featuredImage': 'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiOti36pT-2IwVvWjrezVC87xl1pw2W9SQgYhfI0sRUOSG5EDDtdp6jQR9iJx6k3me_zklHq_2RxaMqYD8628T3bteI0gZ5ZDdogQxSleJppglgBo-MKPDdQUWEEw9vX0UjFvDZ2lU0UT4eOVg64CqmqQPP6AH9BntCVpA5kx3tiQZvrI66_GYxgwwG/s1600/Tensorflow-building-reinforcement-learning-agent-with-JAX-social.png', 'url': 'https://blog.tensorflow.org/2022/09/building-reinforcement-learning-agent-with-JAX-and-deploying-it-on-android-with-tensorflow-lite.html', 'type': 'item', 'isSingleItem': true, 'isMultipleItems': false, 'isError': false, 'isPage': false, 'isPost': true, 'isHomepage': false, 'isArchive': false, 'isLabelSearch': false, 'postId': 6024940900969663155}}]); _WidgetManager._RegisterWidget('_HTMLView', new _WidgetInfo('HTML1', 'nav', document.getElementById('HTML1'), {}, 'displayModeFull')); _WidgetManager._RegisterWidget('_FeaturedPostView', new _WidgetInfo('FeaturedPost1', 'blog', document.getElementById('FeaturedPost1'), {}, 'displayModeFull')); _WidgetManager._RegisterWidget('_BlogView', new _WidgetInfo('Blog1', 'blog', document.getElementById('Blog1'), {'cmtInteractionsEnabled': false, 'lightboxEnabled': true, 'lightboxModuleUrl': 'https://www.blogger.com/static/v1/jsbin/2646514562-lbx.js', 'lightboxCssUrl': 'https://www.blogger.com/static/v1/v-css/1964470060-lightbox_bundle.css'}, 'displayModeFull')); _WidgetManager._RegisterWidget('_HTMLView', new _WidgetInfo('HTML2', 'footer', document.getElementById('HTML2'), {}, 'displayModeFull')); </script> </body> </html>