differential-training/models.py at master · nikhil-dev/differential-training · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import tensorflow as tf

import tensorflow.contrib.eager as tfe

class Actor(tfe.Network):
  """Actor network with two fully connected layers"""
  def __init__(self, nb_actions, nb_obs, layer_norm=True):
	super(Actor, self).__init__(name='actor_scope')
	self.nb_actions = nb_actions
	self.layer_norm = layer_norm
	self._input_shape = [-1, nb_obs]

	self.fc1 = self.track_layer(tf.layers.Dense(64))
	self.fc2 = self.track_layer(tf.layers.Dense(64))
	self.batch_norm = self.track_layer(tf.layers.BatchNormalization())
	self.final_layer = self.track_layer(tf.layers.Dense(self.nb_actions))

  def call(self, obs):
  	x = tf.reshape(obs, self._input_shape)
	x = self.fc1(x)
	if self.layer_norm:
		x = self.batch_norm(x)
	x = tf.nn.relu(x)

	x = self.fc2(x)
	if self.layer_norm:
		x = self.batch_norm(x)
	x = tf.nn.relu(x)

	x = self.final_layer(x)
	return tf.nn.tanh(x)


class Critic(tfe.Network):
  """Critic network with two fully connected layers"""
  def __init__(self, nb_actions, nb_obs, layer_norm=True):
	super(Critic, self).__init__(name='critic_scope')
	self.layer_norm = layer_norm

	self.fc1 = self.track_layer(tf.layers.Dense(64))
	self.fc2 = self.track_layer(tf.layers.Dense(64))
	self.batch_norm = self.track_layer(tf.layers.BatchNormalization())
	self.final_layer = self.track_layer(tf.layers.Dense(1))
	self._input_action_shape = [-1, nb_actions]
	self._input_obs_shape = [-1, nb_obs]

  def call(self, obs, actions):
	obs = tf.reshape(obs, self._input_obs_shape)
	actions = tf.reshape(actions, self._input_action_shape)

	x = obs
	x = self.fc1(x)
	if self.layer_norm:
		x = self.batch_norm(x)
	x = tf.nn.relu(x)

	x = tf.concat([x, actions], axis=-1)
	x = self.fc2(x)
	if self.layer_norm:
		x = self.batch_norm(x)
	x = tf.nn.relu(x)

	return self.final_layer(x)


class DifferenceCritic(tfe.Network):
  def __init__(self, nb_actions, nb_obs, layer_norm=True):
	super(DifferenceCritic, self).__init__(name='difference_critic_scope')
	self.layer_norm = layer_norm

	self.fc1 = self.track_layer(tf.layers.Dense(64))
	self.fc2 = self.track_layer(tf.layers.Dense(64))
	self.batch_norm = self.track_layer(tf.layers.BatchNormalization())
	self.final_layer = self.track_layer(tf.layers.Dense(1))
	self._input_action_shape = [-1, nb_actions]
	self._input_obs_shape = [-1, nb_obs]

  def call(self, obs1, action1, obs2, action2):
	obs1 = tf.reshape(obs1, self._input_obs_shape)
	obs2 = tf.reshape(obs2, self._input_obs_shape)
	action1 = tf.reshape(action1, self._input_action_shape)
	action2 = tf.reshape(action2, self._input_action_shape)

	x = tf.concat([obs1, action1, obs2, action2], axis=-1)

	x = self.fc1(x)
	if self.layer_norm:
		x = self.batch_norm(x)
	x = tf.nn.relu(x)

	x = self.fc2(x)
	if self.layer_norm:
	  x = self.batch_norm(x)
	x = tf.nn.relu(x)

	return self.final_layer(x)