Performance drop when transform codes from Pytorch to JAX

[dui1234]

Hi,

I am trying to code a policy gradient reinforcement learning in JAX by transforming the script provided in https://github.com/tsmatz/reinforcement-learning-tutorials/blob/master/02-policy-gradient.ipynb as following:

import jax.numpy as jnp
import numpy as np
import jax 
from jax import random,lax,vmap,jit
from jax.example_libraries import stax, optimizers
from jax.nn import initializers
from functools import partial
import gym
import optax
import matplotlib.pylab as plt

def get_NN(D_in,D_out,H = 64,seed=5407):
    key = jax.random.PRNGKey(seed)
    key,sk = random.split(key, 2)
    
    Winit = initializers.he_uniform(dtype=jnp.float64)
    binit = initializers.normal(stddev=0.1, dtype=jnp.float64)
    
    net_init, net_forward_fn = stax.serial(
                                stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(D_out, Winit, binit))
    output_shape, net_params = net_init(sk, (D_in,))
        
    return [net_params,net_forward_fn]

def get_agent(s_params):
    policy_NN = get_NN(s_params['obs_d'],s_params['act_d'])
    lr = 0.001
    opt_init, opt_update, get_params = optimizers.adam(step_size=lr)
    opt_state = opt_init(policy_NN[0])
    return opt_state,policy_NN[1], opt_update, get_params

@partial(jit,static_argnums = (0,1,))
def get_action(forward_f,get_params,opt_state,key,s):
    NN_param = get_params(opt_state)
    logits = forward_f(NN_param,s)
    key,sk = random.split(key, 2)
    a = random.categorical(sk,logits)
    return a,key

@partial(jit, static_argnums=(1,))
def cal_loss(NN_param,forward_f,states,actions,cum_rewards):
    logits = forward_f(NN_param,states)
    logits = jax.nn.log_softmax(logits)
    n_classes = logits.shape[-1]
    log_probs = -optax.softmax_cross_entropy(logits,jax.nn.one_hot(actions, n_classes))
    loss = -log_probs * cum_rewards
    return jnp.mean(loss)

@partial(jit, static_argnums=(0,1,2))
def update_model(get_params,forward_f,opt_update,step,opt_state,states,actions,cum_rewards):
    NN_param = get_params(opt_state)
    
    loss, grads = jax.value_and_grad(cal_loss)(NN_param,forward_f,states,actions,cum_rewards)
    return loss, opt_update(step,grads,opt_state)

seed = 999
s_param = {'obs_d':4,'act_d':2}
params_NN,forward_f,opt_update,get_params = get_agent(s_param)
env = gym.make("CartPole-v1")
reward_records = []
losses = []
key = jax.random.PRNGKey(seed)

for i in range(1000):
    s = env.reset()
    done = False
    states = []
    actions = []
    rewards = []
    
    while not done:
        states.append(s.tolist())
        action,key = get_action(forward_f,get_params,params_NN,key,s)
        s, r, term, trunc = env.step(np.asarray(action))
        done = term or trunc
        actions.append(np.asarray(action))
        rewards.append(r)
        
    # Get cumulative rewards
    cum_rewards = np.zeros_like(rewards)
    reward_len = len(rewards)
    for j in reversed(range(reward_len)):
        cum_rewards[j] = rewards[j] + (cum_rewards[j+1]*gamma if j+1 < reward_len else 0)
        
    # Train (optimize parameters)
    states_j = jnp.array(states)
    actions_j = jnp.array(actions)
    cum_rewards_j = jnp.array(cum_rewards)

    loss, opt_state = update_model(get_params,forward_f,opt_update,i,params_NN,states_j,actions_j,cum_rewards_j)
    params_NN = opt_state
    
    print("Run episode{} with rewards {}".format(i, sum(rewards)), end="\r")
    reward_records.append(sum(rewards))
    losses.append(loss)
    
print("\nDone")
env.close()

plt.plot(reward_records)

Compared to the original Pytorch code, the performance is much worse. I suspected that it may be caused by the initialization so I tried changing these lines with available jax initializers options:

Winit = initializers.he_uniform(dtype=jnp.float64) binit = initializers.normal(stddev=0.1, dtype=jnp.float64)

However, I ran into another problem that when I use the binit = initializers.he_uniform(stddev=0.1, dtype=jnp.float64) or something else rather than normal, it gives me an error.

I am very much obliged for you help!!

[jakevdp]

Hi - I'm not precisely sure what might be causing the performance degredation, but one thing stands out: your code makes use of jax.example_libraries: that code is not intended for use in this manner, as mentioned in the docstring: https://github.com/google/jax/blob/d6e06f4476c150b285e605d7346aa06dc4d204dc/jax/example_libraries/stax.py#L17-L20

Code in jax.example_libraries has not been heavily developed or optimized, so it would not surprise me if your program is slower than it could be. If you're interested in performance in JAX, I'd suggest working with tools that have been written with performance and usability in mind: for example, flax is a good choice for defining neural networks.

[dui1234]

Thank you for pointing this out, I have just realized this fact 😹🙏. I will try implementing it with flax to see if I can get the same performance.

[dui1234]

Hi, just want to update you the implementation using Flax. The performance is still inferior compared to the Torch version. Up to this point, I suppose the NN module should not be the casue of the inferiority but I still cannot figure it out what else might causes it. The implementation including Flax is as below:

from jax.config import config; config.update("jax_enable_x64",True)
import jax
import jax.numpy as jnp
import optax
import gym
from functools import partial
import numpy as np
from jax import random as jrand
from collections import deque
from flax import linen as nn
from flax.training.train_state import TrainState
import flax
from typing import Sequence

class NN(nn.Module):
    hiddden: Sequence[int]

    @nn.compact
    def __call__(self,input):
        x = input
        for i, num_neu in enumerate(self.hiddden):
            x = nn.Dense(num_neu)(x)
            if i != len(self.hiddden) - 1:
                x = nn.relu(x)
        return x

def creat_train_state(module, key, obs):
    params = module.init(key, obs)
    tx = optax.adam(learning_rate = 0.001)
    return TrainState.create(
        apply_fn = jax.jit(module.apply), 
        params = params,
        tx = tx
    )

@jax.jit
def get_action(state,obs,key):
    logist = state.apply_fn(state.params, obs)
    probs = nn.softmax(logist)
    key, sk = jrand.split(key)
    a = jrand.choice(sk,jnp.arange(probs.shape[0]),p=probs)
    #a = jrand.categorical(sk,probs)
    return a,key

@jax.jit
def update(state,batch):

    def cal_loss(params):
        logits = nn.log_softmax(state.apply_fn(params, batch['state']))
        n_classes = logits.shape[-1]
        log_probs = -optax.softmax_cross_entropy(logits,nn.one_hot(batch['action'], n_classes))
        return jnp.mean(-log_probs * batch['cum_rew'])
        
    loss, grads = jax.value_and_grad(cal_loss)(state.params)
    state = state.apply_gradients(grads=grads)
    return loss, state

env = gym.make("CartPole-v1")
act_dim = env.action_space.n
key, ini_key = jrand.split(jrand.PRNGKey(0))

module = NN([128,128,act_dim])
s = env.reset()[0]
NN_state = creat_train_state(module, ini_key, s)
gamma = 0.99
reward_records = []
losses = []

for i in range(1000):
    s = env.reset()[0]
    done = False
    states = []
    actions = []
    rewards = [] 

    while not done:
        states.append(s.tolist())
        a, key = get_action(NN_state,s,key)
        s, r, term, trunc, _ = env.step(np.asarray(a))
        done = term or trunc
        actions.append(a)
        rewards.append(r)

    cum_rewards = np.zeros_like(rewards)
    reward_len = len(rewards)
    for j in reversed(range(reward_len)):
        cum_rewards[j] = rewards[j] + (cum_rewards[j+1]*gamma if j+1 < reward_len else 0)

    batch = {'state':jnp.array(states),'action':jnp.array(actions),'cum_rew':jnp.array(cum_rewards)}
    loss, NN_state = update(NN_state, batch)

    print("Run episode {} with rewards {}".format(i, sum(rewards)), end="\r")
    reward_records.append(sum(rewards))
    losses.append(float(loss))

print("\nDone")
env.close()
plt.plot(reward_records)

Your help will be very much obliged !!

[jakevdp]

For flax-specific questions, you might have more luck asking at http://github.com/google/flax.

[ayaka14732]

Why is this line:

params_NN = opt_state

[dui1234]

My working network parameter is param_NN, I need to update it after the weighs and biases update. The output NN parameters of the update is opt_state, and I update the working NN parameters with the line.