fix(ppo): added message_passer to ppo
This commit is contained in:
Robin Meersman
parent
079480324d
commit
af68a0064c
4 changed files with 27 additions and 38 deletions
|
|
@ -40,16 +40,17 @@ class Actor(nn.Module):
|
|||
class MessagePasser(nn.Module):
|
||||
hidden_dim: int
|
||||
num_propagation_steps: int
|
||||
adj_matrix: jnp.ndarray
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x: jnp.ndarray, adj_matrix: jnp.ndarray):
|
||||
def __call__(self, x: jnp.ndarray):
|
||||
for _ in range(self.num_propagation_steps):
|
||||
# (n_nodes, feat)
|
||||
messages = nn.Dense(self.hidden_dim)(x)
|
||||
messages = nn.tanh(messages)
|
||||
|
||||
# note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8)
|
||||
agg = adj_matrix
|
||||
agg = self.adj_matrix
|
||||
aggregated = agg @ messages
|
||||
|
||||
x_concat = jnp.concatenate([x, aggregated], axis=-1)
|
||||
|
|
|
|||
|
|
@ -72,6 +72,10 @@ def create_obs_processor(
|
|||
padded[key] = arr
|
||||
return padded
|
||||
|
||||
# TODO
|
||||
def _split_to_agents() -> dict:
|
||||
return {}
|
||||
|
||||
def _flatten_features(obs: dict) -> jnp.ndarray:
|
||||
ordered_keys = [
|
||||
"disk_z_tilt",
|
||||
|
|
@ -93,7 +97,7 @@ def create_obs_processor(
|
|||
if v.size == 0:
|
||||
continue
|
||||
|
||||
# reshape scalars
|
||||
# reshape scalars to 1D array
|
||||
if v.ndim == 0:
|
||||
v = v.reshape(1)
|
||||
|
||||
|
|
@ -102,19 +106,8 @@ def create_obs_processor(
|
|||
values.append(v.reshape(1, -1)) # (1, feat)
|
||||
continue
|
||||
|
||||
# -------- SPLIT TO SEGMENTS --------
|
||||
# -------- SCALE WITH SEGMENTS --------
|
||||
if key in _JOINT_SCALED_KEYS:
|
||||
if morph_mode == MorphMode.SEGMENT:
|
||||
joint_count = 3
|
||||
axis_per_joint = 2
|
||||
|
||||
center_size = num_arms * joint_count * axis_per_joint
|
||||
v_center = v[:center_size].reshape(
|
||||
num_arms, joint_count * axis_per_joint
|
||||
) # (arms, 6)
|
||||
v_segs = v[center_size:].reshape(-1, 2) # (segs, 2)
|
||||
values.append(jnp.concatenate([v_center, v_segs], axis=0)) # (arms+segs, ?)
|
||||
continue
|
||||
v = v.reshape(num_arms, -1) # (n_arms, 2)
|
||||
|
||||
elif key in _SEGMENT_SCALED_KEYS:
|
||||
|
|
@ -122,25 +115,20 @@ def create_obs_processor(
|
|||
|
||||
else:
|
||||
# global key, broadcast to all nodes
|
||||
n_nodes = (num_segments + num_arms) if morph_mode == MorphMode.SEGMENT else num_arms
|
||||
v = jnp.repeat(v[None, :], n_nodes, axis=0) # (n_nodes, feat)
|
||||
v = jnp.repeat(v[None, :], num_arms, axis=0) # (num_arms, feat)
|
||||
|
||||
# -------- SEGMENT MODE --------
|
||||
if morph_mode == MorphMode.SEGMENT:
|
||||
values.append(v) # (n_nodes, feat)
|
||||
continue
|
||||
|
||||
# -------- ARM MODE --------
|
||||
# RING + FULLY CONNECTED
|
||||
v = v.reshape(num_arms, -1)
|
||||
values.append(v) # (n_arms, feat)
|
||||
|
||||
return jnp.concatenate(values, axis=-1)
|
||||
return jnp.concatenate(values, axis=-1) # (agent, feat)
|
||||
|
||||
def _process_single(obs_dict: dict) -> jnp.ndarray:
|
||||
processed = _add_derived_features(obs_dict)
|
||||
processed = _normalize_features(processed)
|
||||
processed = _add_derived_features(obs_dict) # key |--> (feat-count, feat-lengths, )
|
||||
processed = _normalize_features(processed) # key |--> (feat-count, feat-lengths, )
|
||||
# TODO: split to agents # key |--> (agents, feat-count, feat-lengths)
|
||||
if padding_masks is not None:
|
||||
processed = _pad_features(processed)
|
||||
return _flatten_features(processed)
|
||||
processed = _pad_features(processed) # key |--> (agents, feat-count, feat-lengths')
|
||||
return _flatten_features(processed) # (agents, feat)
|
||||
|
||||
return jax.jit(jax.vmap(_process_single))
|
||||
|
|
|
|||
|
|
@ -119,7 +119,7 @@ def get_action_and_value(
|
|||
):
|
||||
hidden_sensor = sensor_apply(params["sensor_params"], x)
|
||||
hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
|
||||
hidden_sensor = message_passer(hidden_sensor)
|
||||
hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor)
|
||||
|
||||
debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}")
|
||||
debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}")
|
||||
|
|
|
|||
|
|
@ -62,7 +62,6 @@ def _get_action_and_value_noise(
|
|||
key,
|
||||
action_low,
|
||||
action_high,
|
||||
adj_matrix: jnp.ndarray,
|
||||
):
|
||||
# (B, n_nodes, feat)
|
||||
hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
|
||||
|
|
@ -70,7 +69,7 @@ def _get_action_and_value_noise(
|
|||
if message_passer is not None:
|
||||
params = agent_state.params["message_passer_params"]
|
||||
# (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension
|
||||
hidden = jax.vmap(lambda x: message_passer.apply(params, x, adj_matrix))(hidden)
|
||||
hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden)
|
||||
|
||||
hidden_critic = apply_shared(
|
||||
feature_extractor, agent_state.params["feature_extractor_params"], next_obs
|
||||
|
|
@ -100,7 +99,6 @@ def _step_once(
|
|||
carry,
|
||||
_,
|
||||
env_step_fn,
|
||||
adj_matrix,
|
||||
sensor: nn.Module,
|
||||
feature_extractor: nn.Module,
|
||||
actor: nn.Module,
|
||||
|
|
@ -121,7 +119,6 @@ def _step_once(
|
|||
key,
|
||||
action_low,
|
||||
action_high,
|
||||
adj_matrix,
|
||||
)
|
||||
logger11.debug(f"[_step_once] raw_action: {raw_action.shape}")
|
||||
logger11.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}")
|
||||
|
|
@ -242,7 +239,6 @@ def _rollout_jit(
|
|||
message_passer: Optional[nn.Module],
|
||||
action_low,
|
||||
action_high,
|
||||
adj_matrix,
|
||||
):
|
||||
(agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
|
||||
partial(
|
||||
|
|
@ -255,7 +251,6 @@ def _rollout_jit(
|
|||
env_step_fn=step_env_fn,
|
||||
action_low=action_low,
|
||||
action_high=action_high,
|
||||
adj_matrix=adj_matrix,
|
||||
),
|
||||
(agent_state, episode_stats, next_obs, next_done, key, env_state),
|
||||
(),
|
||||
|
|
@ -393,7 +388,6 @@ class PPOTrainer:
|
|||
message_passer=self.message_passer,
|
||||
action_low=action_low,
|
||||
action_high=action_high,
|
||||
adj_matrix=self.adj,
|
||||
)
|
||||
)
|
||||
self._compute_gae_jit = logged_jit(
|
||||
|
|
@ -419,7 +413,13 @@ class PPOTrainer:
|
|||
def apply_feature(p, x):
|
||||
return apply_shared(self.feature_extractor, p, x)
|
||||
|
||||
self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature)
|
||||
def apply_message_passer(p, x):
|
||||
assert self.message_passer is not None
|
||||
return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x)
|
||||
|
||||
self._ppo = PPO(
|
||||
self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature, apply_message_passer
|
||||
)
|
||||
|
||||
self.agent_state = self._init_agent_state()
|
||||
|
||||
|
|
@ -453,6 +453,7 @@ class PPOTrainer:
|
|||
MessagePasser(
|
||||
hidden_dim=300,
|
||||
num_propagation_steps=self.cfg.architecture.message_passing_steps or 4,
|
||||
adj_matrix=self.adj,
|
||||
)
|
||||
if self.morph_mode != MorphMode.CENTRALIZED
|
||||
else None
|
||||
|
|
@ -517,7 +518,6 @@ class PPOTrainer:
|
|||
message_passer_params = self.message_passer.init(
|
||||
message_passer_key,
|
||||
self.sensor.apply(single_sensor_param, sample_obs),
|
||||
self.adj,
|
||||
)
|
||||
self.logger.debug(
|
||||
f"[_init_agent_state] message_passer_params: {
|
||||
|
|
|
|||
Reference in a new issue