<jupyter_start><jupyter_text>testing part<jupyter_code>%%time
## testing results
encode_vector = encoder_one_hot_action(0)
encoder_decode_action(encode_vector)

encoded, encoded_hot = encoder_one_hot_state(obs)
encoder_decode_state(encoded_hot)

## computer result
for i in range(10):
    list_obs, actions = compute_random_suffle(0)<jupyter_output>CPU times: user 227 ms, sys: 50.3 ms, total: 278 ms
Wall time: 275 ms<jupyter_text>Generation part<jupyter_code>def flattenx(x):
    return jnp.ravel(x)

flatten_all = vmap(flattenx)

def generate_dataset(nb_iter_main_loop=100):

    actions_all = []
    obs_all = []

    for _ in range(nb_iter_main_loop):

        # train on main loop
        list_obs, actions = compute_random_suffle(0)

        # process for training
        actions_all.append(actions)
        obs_all.append(list_obs)
        
    obs_all = jnp.concatenate(obs_all)
    actions_all = jnp.concatenate(actions_all)
    
    obs_all, second = encoder_one_hot_state_vector(obs_all)
    actions_all = encoder_one_hot_action_action(actions_all)
        
    return obs_all, actions_all
%%time
obs_all, actions_all = generate_dataset()

dataset = TensorDataset(torch.from_numpy(np.array(obs_all)), torch.from_numpy(np.array(actions_all)))
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

optax.__version__<jupyter_output><empty_output><jupyter_text>Optax optimization<jupyter_code>%%time

from flax.training import train_state

loss_hist = []

optimizer = optax.adam(0.001)

nb_epoch = 1
nb_round_generation = 50
size_generation = 10000

# label smoothing
smooth_vec = vmap(optax.smooth_labels, in_axes=(0, None))

state = train_state.TrainState.create(
      apply_fn=model_rubik.apply, params=params, tx=optimizer)


@jax.jit
def train_step(state, batch):
    """Train for a single step."""
    def loss_fn(params):
        logits = model_rubik.apply({'params': params}, batch['rubik'])
        actions_ = smooth_vec(batch['action'], 0.1)
        loss = jnp.mean(
            optax.softmax_cross_entropy(
                logits=logits, labels=actions_))
        return loss, logits
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (_, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    metrics = compute_metrics(logits=logits, labels=batch['action'])
    return state, metrics

def compute_metrics(logits, labels):
    loss = jnp.mean(
      optax.softmax_cross_entropy(
          logits=logits, labels=labels))
    accuracy = jnp.mean(jnp.argmax(logits, -1) == jnp.argmax(labels, -1))
    metrics = {
      'loss': loss,
      'accuracy': accuracy,
      }
    return metrics

%%time

loss_hist = []
i = 0

for q in range(nb_round_generation):

    #print("New generation")
    
    obs_all, actions_all = generate_dataset(size_generation)

    dataset = TensorDataset(torch.from_numpy(np.array(obs_all)), torch.from_numpy(np.array(actions_all)))
    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

    for p in range(nb_epoch):
        
        print("epoch")
         
        for obs, actions in dataloader:
            
            obs = jnp.array(obs)
            actions = jnp.array(actions)
            
            batch = {}
            batch['action'] = actions
            batch['rubik'] = obs
            
            # optimizer application
            state, metrics = train_step(state, batch)
            
            loss_hist.append(metrics['loss'].item())
            #params = jax.lax.stop_gradient(params)
            i += 1
            
            if (i % 100) == 0:
                #print(i)
                xla._xla_callable.cache_clear()
                
        #jax.profiler.save_device_memory_profile(f"memory{p}_{q}.prof")


plt.plot(loss_hist)
for obs, actions in dataloader:
    obs = jnp.array(obs)
    actions = jnp.array(actions)
loss_hist_conv = np.convolve(loss_hist, np.ones(50)/50)
plt.plot(loss_hist_conv)
loss_hist_conv = np.convolve(loss_hist, np.ones(50)/50)
plt.plot(loss_hist_conv)<jupyter_output><empty_output><jupyter_text>Testing performance for solving rubiks !!!!<jupyter_code>env.reset()

state_goal = env.state.copy()
## test to solve a rubiks
obs, action = compute_random_suffle(0, nb_iter=10)
begin_state = obs[-1, :, :, :]
env.reset()
env.set_init(np.array(begin_state))
state_obs = begin_state
flat_obs, noflat_obs = encoder_one_hot_state(state_obs)


count = 0
while np.all(state_goal == np.array(state_obs)) == False and count < 50:
    
    action = jnp.argmax(model_rubik.apply({'params': state.params}, flat_obs)).item()
    print(action)
    state_obs, _, _, _= env.step(action)
    flat_obs, noflat_obs = encoder_one_hot_state(state_obs)
    count += 1
    
    
    
state_obs
begin_state
model_rubik.apply({'params': state.params}, flat_obs)
env.reset()
episode_rew = 0

## action list
#actions = jrandom.randint(key1, shape=(nb_iter,1), minval=0, maxval=12)
actions = np.random.randint(0, 12, size=(nb_iter,1))

## state list
obs_list = []

for q in range(nb_iter):

    print(actions[q, 0])
    obs, rew, done, _ = env.step(actions[q, 0])
    obs_list.append(obs.copy())
    print(obs)

    ## we should also compute the max of the model

### we compute the reverse action
np.array(state_obs)
obs

pickle.dump(state.params, open( "save_state.p", "wb" ) )
params = pickle.load(open( "save_state.p", "rb" ))
params<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup