acnet.py

import tensorflow as tf
import tensorflow.keras.layers as kl
import tensorflow_probability as tfp
import numpy as np


class ActorCriticNet(tf.keras.Model):

    def __init__(self, action_space=2):

        super(ActorCriticNet, self).__init__()

        self.action_space = action_space

        self.dense1 = kl.Dense(100, activation="relu")

        self.dense2 = kl.Dense(100, activation="relu")

        self.values = kl.Dense(1)

        self.policy_logits = kl.Dense(action_space)

    def call(self, x):

        x1 = self.dense1(x)
        logits = self.policy_logits(x1)

        x2 = self.dense2(x)
        values = self.values(x2)

        return values, logits

    def sample_action(self, state):

        state = tf.convert_to_tensor(np.atleast_2d(state), dtype=tf.float32)

        _, logits = self(state)

        action_probs = tf.nn.softmax(logits)

        cdist = tfp.distributions.Categorical(probs=action_probs)

        action = cdist.sample()

        return action.numpy()[0]