DQN Agent Issue With Custom Environment

Viewed 486

I tried implementing my own DQN following the 1_dqn_tutorial from official tensorflow docs. My observation_specs are of the shape (20,20) and action_specs are (3,). The problem I encounter is that whenever I try to create a new agent I get the following error:

ValueError: Expected q_network to emit a floating point tensor with inner dims (3,); but saw network output spec: TensorSpec(shape=(20, 3), dtype=tf.float32, name=None) In call to configurable 'DqnAgent' (<class 'tf_agents.agents.dqn.dqn_agent.DqnAgent'>)

My custom environment:

class LearningEnvironment(py_environment.PyEnvironment):

def __init__(self):
    super().__init__()
    self._action_spec = array_spec.BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=2, name='action')
    self._observation_spec = array_spec.BoundedArraySpec(shape=(BOARD_SIZE, BOARD_SIZE), dtype=np.int32, minimum=0,
                                                         maximum=3, name='observation')

    self._board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=np.int32)
    self._snake_length = 1
    self._episode_ended = False

def observation_spec(self):
    return self._observation_spec

def action_spec(self):
    return self._action_spec

def _step(self, action):

    if self._episode_ended:
        return self.reset()

    if action == 0:
        res = main.next_tick(pyglet.window.key.A)
        if res is None:
            self._episode_ended = True
        else:
            self._board = np.array(res[0], dtype=np.int32)
            self._snake_length = res[1]
    elif action == 1:
        res = main.next_tick(pyglet.window.key.D)
        if res is None:
            self._episode_ended = True
        else:
            self._board = np.array(res[0], dtype=np.int32)
            self._snake_length = res[1]
    elif action == 2:
        res = main.next_tick(0)
        if res is None:
            self._episode_ended = True
        else:
            self._board = np.array(res[0], dtype=np.int32)
            self._snake_length = res[1]
    else:
        raise ValueError('action should be 0, 1 or 2')

    if self._episode_ended:
        reward = self._snake_length
        return time_step.termination(self._board, reward)
    else:
        return time_step.transition(self._board, reward=0.0, discount=1.0)

def _reset(self):
    self._snake_length = 1
    self._episode_ended = False
    return time_step.restart(self._board)

My DQN attempt:

fc_layer_params = (100, 50)
action_tensor_spec = tensor_spec.from_spec(env.action_spec())
num_actions = action_tensor_spec.maximum - action_tensor_spec.minimum + 1

def dense_layer(num_inputs):
    return tf.keras.layers.Dense(
        num_inputs,
        activation=tf.keras.activations.relu,
        kernel_initializer=tf.keras.initializers.VarianceScaling(
            scale=2.0, mode='fan_in', distribution='truncated_normal'))


dense_layers = [dense_layer(num_units) for num_units in fc_layer_params]
q_values_layer = tf.keras.layers.Dense(
    num_actions,
    activation=None,
    kernel_initializer=tf.keras.initializers.RandomUniform(
        minval=-0.03, maxval=0.03),
    bias_initializer=tf.keras.initializers.Constant(-0.2))

q_net = sequential.Sequential(dense_layers + [q_values_layer])

optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

train_step_counter = tf.Variable(0)

#print(env.action_spec().shape)

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    optimizer=optimizer,
    td_errors_loss_fn=common.element_wise_squared_loss,
    train_step_counter=train_step_counter)

agent.initialize()
0 Answers
Related