vmap on list?

[dui1234]

Hi,

I am aware that vmap requires an array to be passed into a transformmed function but once I try to use both numpy array or jnp.array I got errors. Here is the code:

import numpy as np
import jax
import jax.numpy as jnp
from jax import random,lax,vmap,jit
from jax.nn import initializers
from jax.example_libraries import stax, optimizers

def get_agent(s_params):
    policy_NN = get_NN(s_params['obs_dim'],s_params['H'],s_params['action_dim'])
    target_NN = get_NN(s_params['obs_dim'],s_params['H'],s_params['action_dim'],seed=10)
    target_NN[0] = policy_NN[0].copy()
    
    lr = 1e-3
    opt_init, opt_update, get_params = optimizers.adam(step_size=lr)
    opt_state = opt_init(policy_NN[0])
    
    return dict(opt_state=opt_state,target_NN=target_NN[0]),policy_NN[1], opt_update, get_params

def get_NN(D_in,H,D_out,seed=5407):
    key = jax.random.PRNGKey(seed)
    key, sk1,sk2 = random.split(key, 3)
    Winit = initializers.lecun_normal(dtype=jnp.float64)
    binit = initializers.normal(stddev=1.0, dtype=jnp.float64)
        
    net_init, net_forward_fn = stax.serial(stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(D_out, Winit, binit))
    output_shape, net_params = net_init(sk2, (D_in,))
        
    return [net_params,net_forward_fn]

def get_act(s_param,sp,NN_param,state):
    return jnp.where(s_param['epsilon'] > np.random.random(), jax.random.randint(sp, shape=(1,), minval=0, 
                                                         maxval=s_param['action_dim'], 
                                                         dtype=jnp.uint8),
                               forward_f(NN_param,state).argmax())

system_params = {'epsilon': 1.0,'obs_dim':4,'action_dim':2,'H':128,'N_swim':2}

I, then, tried to construct a list contain dictionaries, each of them represent an agent in my system and one agent carries two networks:

agents = []
for i in range(system_params['N_swim']):
    agent,forward_f,opt_update,get_params = get_agent(system_params)
    agents.append(agent)

Then the network parameters of each agent are extracted to a list as:

params1 = []
for i in range(len(agents)):
    params1.append(get_params(agents[i]['opt_state']))

Get a dummy set of information to be passed to each of the agent's networks:
obs = jnp.array([np.random.randint(10,size=4) for i in range(2)])

pass all information to a function in order to get an action that each agent will perform for the next time step:

seed = 111
key = random.PRNGKey(seed)
keys = random.split(key,system_params['N_swim']+1)
key, sp = keys[0],keys[1:]

get_act(system_params,sp[0],params1[0],obs[0])

This works perfectly fine with the result of: DeviceArray([0], dtype=int32)

However, when I tried to vmap the get_act function and pass the whole batch of information as:

vmap(get_act,(None,0,0,0))(system_params,sp,params1,obs)
This gives me error due to the fact that the params1 is a list. I, then, tried to converte the params1 to a numpy array as pass to the vmapped-function again:

params1_np = np.array(params1)
vmap(get_act,(None,0,0,0))(system_params,sp,params1_np,obs)

This raises an AssertionError:

When I try to convert the params1 to jnp.array, it is also prohibited as the shapes of each leave are not equal:

params1_jnp = jnp.array(params1) . . . ValueError: All input arrays must have the same shape.

Could you help me out of this? 😂😂

[soraros]

Could you provide a runnable example (preferably in a single code block)? I don't see params1 defined anywhere.

[dui1234]

So sorry, I missed a cell which contatins the params1. Just updated the thread.

[jakevdp]

It's unclear to me what forward_f in your function is... is it meant to just be the value from the last iteration of constructing agents?

I appreciate the narrative explaining the reasoning in your code, but I'd suggest spending a few minutes making a minimal reproducible example – i.e. a single block of code that someone could paste into a fresh CPU Colab runtime, and execute without modification to see the error you're seeing.

[dui1234]

The forward_f is a function that obtain as an return when calling stax.serial and yes, it is the last one obtain from the last iteration.

Here is the one-cell block version of my code:

import numpy as np
import jax
import jax.numpy as jnp
from jax import random,lax,vmap,jit
from jax.nn import initializers
from jax.example_libraries import stax, optimizers

def get_agent(s_params):
    policy_NN = get_NN(s_params['obs_dim'],s_params['H'],s_params['action_dim'])
    target_NN = get_NN(s_params['obs_dim'],s_params['H'],s_params['action_dim'],seed=10)
    target_NN[0] = policy_NN[0].copy()
    
    lr = 1e-3
    opt_init, opt_update, get_params = optimizers.adam(step_size=lr)
    opt_state = opt_init(policy_NN[0])
    
    return dict(opt_state=opt_state,target_NN=target_NN[0]),policy_NN[1], opt_update, get_params

def get_NN(D_in,H,D_out,seed=5407):
    key = jax.random.PRNGKey(seed)
    key, sk1,sk2 = random.split(key, 3)
    Winit = initializers.lecun_normal(dtype=jnp.float64)
    binit = initializers.normal(stddev=1.0, dtype=jnp.float64)
        
    net_init, net_forward_fn = stax.serial(stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(D_out, Winit, binit))
    output_shape, net_params = net_init(sk2, (D_in,))
        
    return [net_params,net_forward_fn]

system_params = {'epsilon': 1.0,'obs_dim':4,'action_dim':2,'H':128,'N_swim':2}

agents = []
for i in range(system_params['N_swim']):
    agent,forward_f,opt_update,get_params = get_agent(system_params)
    agents.append(agent)
    
params1 = []
for i in range(len(agents)):
    params1.append(get_params(agents[i]['opt_state']))
    
def get_act(s_param,sp,NN_param,state):
    return jnp.where(s_param['epsilon'] > np.random.random(), jax.random.randint(sp, shape=(1,), minval=0, 
                                                         maxval=s_param['action_dim'], 
                                                         dtype=jnp.uint8),
                               forward_f(NN_param,state).argmax())
    
seed = 111
key = random.PRNGKey(seed)
keys = random.split(key,system_params['N_swim']+1)
key, sp = keys[0],keys[1:]
obs = jnp.array([np.random.randint(10,size=4) for i in range(2)])

get_act(system_params,sp[0],params1[1],obs[0])
####This is ok

#vmap(get_act,(None,0,0,0))(system_params,sp,params1,obs)
####The above line will give an error as the params1 is a list

#params1_jnp = jnp.array(params1)
#vmap(get_act,(None,0,0,0))(system_params,sp,params1_jnp,obs)
####This will gives ValueError: All input arrays must have the same shape.

params1_np = np.array(params1)
vmap(get_act,(None,0,0,0))(system_params,sp,params1_np,obs)

I commented out the other two cases that also produce error. Hope this helps.

[jakevdp]

You cannot use a simple jnp.array(params1) because params1 is not a simple list of arrays, it is a nontrivial nested data structure:

print(jax.tree_map(np.shape, params1))

[[((4, 128), (128,)), (), ((128, 128), (128,)), (), ((128, 2), (2,))],
 [((4, 128), (128,)), (), ((128, 128), (128,)), (), ((128, 2), (2,))]]

As you found, np.array(params1) returns an output, but it is an array with dtype=object, which JAX doesn't support. Numpy actually warns you about this issue as well:

_ = np.array(params1)

VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of
  lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you
  must specify 'dtype=object' when creating the ndarray.

If you want to vmap over your parameters, you need to use a struct-of-arrays rather than an array-of-structs pattern. In your case you can do that with something like this:

params1_stacked = jax.tree_map(lambda *x: jnp.stack(x), *params1)
out = vmap(get_act,(None,0,0,0))(system_params,sp,params1_stacked,obs)
print(out)

[[0]
 [0]]

Here params1_stacked has the same tree structure (i.e. structure of nested lists and tuples) as params1[0], but in place of each individual parameter array is an array with all parameters stacked along the leading dimension. So params_stacked is a struct-of-arrays version of params1, which is why it works with vmap.

One side-note however: as written, your function likely will not work as expected within vmap, because you're calling np.random.random() which is evaluated once for all batches. If you want a different random number for each batch, you should use jax.random with a vmapped key.

[dui1234]

Thank you for the instructive reply, very much obliged!!