I recreated some code I found online for solving the bandits problem using policy gradient. The example was in tensorflow 1.0 so I recreated it with tensorflow 2.0 using eager execution and gradient tape, however, when training the model, I have to convert the weights Tensor to a numpy array, update the weights then reassign the tf.Variable back from the numpy array. I feel this is not performant and could find a better way. Full code is here https://github.com/entrpn/reinforcement-learning/blob/master/tf2_rl/bandits.py

The main code I'm looking to improve as follows:

def train(agent,action,reward, learning_rate=0.001):
    with tf.GradientTape() as t:
        current_loss = loss(agent(action),reward)
    dW = t.gradient(current_loss,[agent.weights])
    weights_as_np = agent.weights.numpy()
    responsible_weight = agent.weights[action]
    responsible_weight_dw = np.array(dW)[0][action]

    weights_as_np[action] = weights_as_np[action] - learning_rate*responsible_weight_dw



1 Answer 1


If possible, try to use just tensorflow functions (put the .numpy() part out, for example), and put a @tf.function decorator on top of train() function.

The role of @tf.function is that it transforms a whole function into a tensorflow op. The whole function will be execute an order of magnitude faster than a normal Python function.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.