Stefan_Haerter:

and now I want to write a custom policy as well.


You could also write an “MVP” using, for example, an SAC 3 (by BAIR) algorithm with just TF Probability and TF as a start.

Here’s a “clean” example:

  

      github.com
  

  
    philtabor/Youtube-Code-Repository/blob/master/ReinforcementLearning/PolicyGradient/SAC/tf2/networks.py 3


    import os
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow_probability as tfp
from tensorflow.keras.layers import Dense 

class CriticNetwork(keras.Model):
    def __init__(self, n_actions, fc1_dims=256, fc2_dims=256,
            name='critic', chkpt_dir='tmp/sac'):
        super(CriticNetwork, self).__init__()
        self.fc1_dims = fc1_dims
        self.fc2_dims = fc2_dims
        self.n_actions = n_actions
        self.model_name = name
        self.checkpoint_dir = chkpt_dir
        self.checkpoint_file = os.path.join(self.checkpoint_dir, name+'_sac')

        self.fc1 = Dense(self.fc1_dims, activation='relu')
        self.fc2 = Dense(self.fc2_dims, activation='relu')



  This file has been truncated. show original

  

  
    
    
  

  


(Another example with code: Deep Reinforcement Learning: Playing CartPole through Asynchronous Advantage Actor Critic (A3C) with tf.keras and eager execution — The TensorFlow Blog 3)