fix: clamp std before distribution, lower init_noise to 0.5, NaN guard
This commit is contained in:
8x54zj-m
13
motrix_rl/src/motrix_rl/dwaq_rsl/__init__.py
Normal file
13
motrix_rl/src/motrix_rl/dwaq_rsl/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
||||
"""Faithful port of Manaro-Alpha/DreamWaQ's custom rsl_rl-1.0.2 for MotrixLab.
|
||||
|
||||
Provides ActorCritic_DWAQ (CENet+Actor+Critic), PPO (with joint VAE training),
|
||||
RolloutStorage, OnPolicyRunner, and a MotrixLab NpEnv -> VecEnv adapter.
|
||||
"""
|
||||
from .actor_critic_dwaq import ActorCritic_DWAQ
|
||||
from .ppo import PPO
|
||||
from .rollout_storage import RolloutStorage
|
||||
from .runner import OnPolicyRunner
|
||||
from .env_wrapper import DwaqVecEnv
|
||||
|
||||
__all__ = ["ActorCritic_DWAQ", "PPO", "RolloutStorage", "OnPolicyRunner", "DwaqVecEnv"]
|
||||
169
motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py
Normal file
169
motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py
Normal file
@@ -0,0 +1,169 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
|
||||
class ActorCritic_DWAQ(nn.Module):
|
||||
def __init__(self, num_actor_obs, num_critic_obs, num_actions, cenet_in_dim, cenet_out_dim, activation="elu", init_noise_std=1.0,):
|
||||
super().__init__()
|
||||
|
||||
self.activation = get_activation(activation)
|
||||
actor_input_dim = num_actor_obs
|
||||
critic_input_dim = num_critic_obs
|
||||
|
||||
self.actor = nn.Sequential(
|
||||
nn.Linear(actor_input_dim,512),
|
||||
self.activation,
|
||||
nn.Linear(512,256),
|
||||
self.activation,
|
||||
nn.Linear(256,128),
|
||||
self.activation,
|
||||
nn.Linear(128,num_actions)
|
||||
)
|
||||
|
||||
self.critic = nn.Sequential(
|
||||
nn.Linear(critic_input_dim,512),
|
||||
self.activation,
|
||||
nn.Linear(512,256),
|
||||
self.activation,
|
||||
nn.Linear(256,128),
|
||||
self.activation,
|
||||
nn.Linear(128,1)
|
||||
)
|
||||
|
||||
self.encoder = nn.Sequential(
|
||||
nn.Linear(cenet_in_dim,128),
|
||||
self.activation,
|
||||
nn.Linear(128,64),
|
||||
self.activation,
|
||||
)
|
||||
self.encode_mean_latent = nn.Linear(64,cenet_out_dim-3)
|
||||
self.encode_logvar_latent = nn.Linear(64,cenet_out_dim-3)
|
||||
self.encode_mean_vel = nn.Linear(64,3)
|
||||
self.encode_logvar_vel = nn.Linear(64,3)
|
||||
|
||||
self.decoder = nn.Sequential(
|
||||
nn.Linear(cenet_out_dim,64),
|
||||
self.activation,
|
||||
nn.Linear(64,128),
|
||||
self.activation,
|
||||
nn.Linear(128,45)
|
||||
)
|
||||
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
# disable args validation for speedup
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
# seems that we get better performance without init
|
||||
# self.init_memory_weights(self.memory_a, 0.001, 0.)
|
||||
# self.init_memory_weights(self.memory_c, 0.001, 0.)
|
||||
|
||||
@staticmethod
|
||||
# not used at the moment
|
||||
def init_weights(sequential, scales):
|
||||
[
|
||||
torch.nn.init.orthogonal_(module.weight, gain=scales[idx])
|
||||
for idx, module in enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))
|
||||
]
|
||||
|
||||
def reset(self, dones=None):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def reparameterise(self,mean,logvar):
|
||||
var = torch.exp(logvar*0.5)
|
||||
code_temp = torch.randn_like(var)
|
||||
code = mean + var*code_temp
|
||||
return code
|
||||
|
||||
def cenet_forward(self,obs_history):
|
||||
distribution = self.encoder(obs_history)
|
||||
mean_latent = self.encode_mean_latent(distribution)
|
||||
logvar_latent = self.encode_logvar_latent(distribution)
|
||||
# var = torch.exp(logvar_latent*0.5)
|
||||
# code_temp = torch.randn_like(var)
|
||||
# code = mean_latent + var*code_temp
|
||||
# print("latent : ",code[0])
|
||||
mean_vel = self.encode_mean_vel(distribution)
|
||||
logvar_vel = self.encode_mean_vel(distribution)
|
||||
code_latent = self.reparameterise(mean_latent,logvar_latent)
|
||||
code_vel = self.reparameterise(mean_vel,logvar_vel)
|
||||
code = torch.cat((code_vel,code_latent),dim=-1)
|
||||
decode = self.decoder(code)
|
||||
return code,code_vel,decode,mean_vel,logvar_vel,mean_latent,logvar_latent
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, observations):
|
||||
mean = self.actor(observations)
|
||||
# Clamp raw actor output before building the Normal distribution:
|
||||
# on unstable terrain (hfield) the network can produce extreme values
|
||||
# that overflow → NaN → crash the full PPO pipeline.
|
||||
mean = torch.clamp(mean, -10.0, 10.0)
|
||||
self.distribution = Normal(mean, mean * 0.0 + self.std)
|
||||
|
||||
def act(self, observations, obs_history, **kwargs):
|
||||
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
|
||||
observations = torch.cat((code,observations),dim=-1)
|
||||
self.update_distribution(observations)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, observations,obs_history):
|
||||
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
|
||||
observations = torch.cat((code,observations),dim=-1)
|
||||
actions_mean = self.actor(observations)
|
||||
return actions_mean
|
||||
|
||||
def evaluate(self, critic_observations, **kwargs):
|
||||
value = self.critic(critic_observations)
|
||||
return value
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.CReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
106
motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py
Normal file
106
motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py
Normal file
@@ -0,0 +1,106 @@
|
||||
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
||||
"""Adapter: MotrixLab DreamWaQ NpEnv -> upstream rsl_rl-1.0.2 VecEnv interface.
|
||||
|
||||
Upstream OnPolicyRunner expects env.step() to return a 7-tuple:
|
||||
(obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, extras)
|
||||
and env.reset() to return (obs, privileged_obs, prev_privileged_obs, obs_hist).
|
||||
|
||||
MotrixLab's DreamWaQTask (numpy, CPU physics) already provides:
|
||||
state.obs -> obs (N, 45)
|
||||
info["privileged_obs"] -> privileged (N, 235) = obs(45)+base_vel(3)+heights(187)
|
||||
info["obs_history"] -> history (N, 5, 45) (flattened to 225)
|
||||
state.terminated/truncated -> dones / time_outs
|
||||
|
||||
This wrapper does numpy<->torch transfer and maintains prev_privileged_obs.
|
||||
"""
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class DwaqVecEnv:
|
||||
"""Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner."""
|
||||
|
||||
def __init__(self, env, device, num_obs=45, num_privileged_obs=235,
|
||||
num_obs_hist=5, num_actions=12, clip_actions=100.0):
|
||||
self._env = env
|
||||
self.device = device
|
||||
self.num_envs = env.num_envs
|
||||
self.num_obs = num_obs
|
||||
self.num_privileged_obs = num_privileged_obs
|
||||
self.num_obs_hist = num_obs_hist
|
||||
self.num_actions = num_actions
|
||||
self._clip_actions = clip_actions
|
||||
self.max_episode_length = int(getattr(env.cfg, "max_episode_steps", None) or 1000)
|
||||
|
||||
n = self.num_envs
|
||||
self.obs_buf = torch.zeros(n, num_obs, device=device)
|
||||
self.obs_hist_buf = torch.zeros(n, num_obs_hist * num_obs, device=device)
|
||||
self.privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
|
||||
self.prev_privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
|
||||
self.rew_buf = torch.zeros(n, device=device)
|
||||
self.reset_buf = torch.ones(n, device=device, dtype=torch.long)
|
||||
self.episode_length_buf = torch.zeros(n, device=device, dtype=torch.long)
|
||||
self.extras = {}
|
||||
|
||||
self._env.init_state()
|
||||
self._sync_from_state(self._env._state)
|
||||
|
||||
# ── helpers ──
|
||||
def _t(self, arr):
|
||||
return torch.from_numpy(np.ascontiguousarray(arr, dtype=np.float32)).to(self.device)
|
||||
|
||||
def _sync_from_state(self, state):
|
||||
self.obs_buf = self._t(state.obs)
|
||||
info = state.info
|
||||
priv = info.get("privileged_obs")
|
||||
if priv is None:
|
||||
priv = np.zeros((self.num_envs, self.num_privileged_obs), dtype=np.float32)
|
||||
self.privileged_obs_buf = self._t(priv)
|
||||
hist = info.get("obs_history") # (N, 5, 45)
|
||||
if hist is None:
|
||||
hist = np.zeros((self.num_envs, self.num_obs_hist, self.num_obs), dtype=np.float32)
|
||||
self.obs_hist_buf = self._t(hist).reshape(self.num_envs, -1)
|
||||
self.episode_length_buf = torch.from_numpy(
|
||||
np.asarray(info.get("steps", np.zeros(self.num_envs)), dtype=np.int64)).to(self.device)
|
||||
|
||||
# ── VecEnv interface ──
|
||||
def step(self, actions):
|
||||
actions = torch.clip(actions, -self._clip_actions, self._clip_actions)
|
||||
actions_np = actions.detach().cpu().numpy().astype(np.float32)
|
||||
# prev_privileged = privileged BEFORE this step (matches upstream)
|
||||
self.prev_privileged_obs_buf = self.privileged_obs_buf.clone()
|
||||
state = self._env.step(actions_np)
|
||||
self._sync_from_state(state)
|
||||
|
||||
done = np.logical_or(state.terminated, state.truncated)
|
||||
self.rew_buf = self._t(state.reward).reshape(-1)
|
||||
self.reset_buf = torch.from_numpy(done.astype(np.int64)).to(self.device)
|
||||
self.extras = {
|
||||
"time_outs": torch.from_numpy(np.asarray(state.truncated, dtype=np.float32)).to(self.device),
|
||||
}
|
||||
# pass per-term reward logging through if present
|
||||
terms = state.info.get("reward_terms")
|
||||
if terms:
|
||||
self.extras["episode"] = {f"rew_{k}": float(v) for k, v in terms.items()}
|
||||
# Log terrain_level so tensorboard tracks curriculum progress.
|
||||
tl = state.info.get("terrain_level")
|
||||
if tl is not None:
|
||||
self.extras["episode"] = self.extras.get("episode", {})
|
||||
self.extras["episode"]["terrain_level"] = float(np.asarray(tl).mean())
|
||||
|
||||
# rewards/dones are (N,) like upstream (storage views them to (N,1) internally)
|
||||
return (self.obs_buf, self.privileged_obs_buf, self.prev_privileged_obs_buf,
|
||||
self.obs_hist_buf, self.rew_buf, self.reset_buf, self.extras)
|
||||
|
||||
def reset(self, env_ids=None):
|
||||
self._env.init_state()
|
||||
self._sync_from_state(self._env._state)
|
||||
self.prev_privileged_obs_buf = torch.zeros_like(self.privileged_obs_buf)
|
||||
return (self.obs_buf, self.privileged_obs_buf,
|
||||
self.prev_privileged_obs_buf, self.obs_hist_buf)
|
||||
|
||||
def get_observations(self):
|
||||
return self.obs_buf, self.obs_hist_buf
|
||||
|
||||
def get_privileged_observations(self):
|
||||
return self.privileged_obs_buf, self.prev_privileged_obs_buf
|
||||
203
motrix_rl/src/motrix_rl/dwaq_rsl/ppo.py
Normal file
203
motrix_rl/src/motrix_rl/dwaq_rsl/ppo.py
Normal file
@@ -0,0 +1,203 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from .actor_critic_dwaq import ActorCritic_DWAQ
|
||||
from .rollout_storage import RolloutStorage
|
||||
|
||||
class PPO:
|
||||
actor_critic: ActorCritic_DWAQ
|
||||
def __init__(self,
|
||||
actor_critic,
|
||||
num_learning_epochs=1,
|
||||
num_mini_batches=1,
|
||||
clip_param=0.2,
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
value_loss_coef=1.0,
|
||||
entropy_coef=0.0,
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
device='cpu',
|
||||
):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
|
||||
# PPO components
|
||||
self.actor_critic = actor_critic
|
||||
self.actor_critic.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
|
||||
self.transition = RolloutStorage.Transition()
|
||||
|
||||
# PPO parameters
|
||||
self.clip_param = clip_param
|
||||
self.num_learning_epochs = num_learning_epochs
|
||||
self.num_mini_batches = num_mini_batches
|
||||
self.value_loss_coef = value_loss_coef
|
||||
self.entropy_coef = entropy_coef
|
||||
self.gamma = gamma
|
||||
self.lam = lam
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.use_clipped_value_loss = use_clipped_value_loss
|
||||
|
||||
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, obs_hist_shape, action_shape):
|
||||
self.storage = RolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, obs_hist_shape, action_shape, self.device)
|
||||
|
||||
def test_mode(self):
|
||||
self.actor_critic.test()
|
||||
|
||||
def train_mode(self):
|
||||
self.actor_critic.train()
|
||||
|
||||
def act(self, obs, critic_obs, prev_critic_obs, obs_history):
|
||||
# if self.actor_critic.is_recurrent:
|
||||
# self.transition.hidden_states = self.actor_critic.get_hidden_states()
|
||||
# Compute the actions and values
|
||||
self.transition.actions = self.actor_critic.act(obs,obs_history).detach()
|
||||
self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
|
||||
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
|
||||
self.transition.action_mean = self.actor_critic.action_mean.detach()
|
||||
self.transition.action_sigma = self.actor_critic.action_std.detach()
|
||||
# need to record obs and critic_obs before env.step()
|
||||
self.transition.observations = obs
|
||||
self.transition.observation_history = obs_history
|
||||
self.transition.critic_observations = critic_obs
|
||||
self.transition.prev_critic_obs = prev_critic_obs
|
||||
return self.transition.actions
|
||||
|
||||
def process_env_step(self, rewards, dones, infos):
|
||||
self.transition.rewards = rewards.clone()
|
||||
self.transition.dones = dones
|
||||
# Bootstrapping on time outs
|
||||
if 'time_outs' in infos:
|
||||
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transitions(self.transition)
|
||||
self.transition.clear()
|
||||
self.actor_critic.reset(dones)
|
||||
|
||||
def compute_returns(self, last_critic_obs):
|
||||
last_values= self.actor_critic.evaluate(last_critic_obs).detach()
|
||||
self.storage.compute_returns(last_values, self.gamma, self.lam)
|
||||
|
||||
def update(self,beta=1):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_autoenc_loss = 0
|
||||
# if self.actor_critic.is_recurrent:
|
||||
# generator = self.storage.reccurent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
# else:
|
||||
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
for obs_batch, critic_obs_batch, prev_critic_obs_batch, obs_hist_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
|
||||
old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch in generator:
|
||||
|
||||
|
||||
self.actor_critic.act(obs_batch, obs_hist_batch, masks=masks_batch, hidden_states=hid_states_batch[0])
|
||||
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.actor_critic.evaluate(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1])
|
||||
mu_batch = self.actor_critic.action_mean
|
||||
sigma_batch = self.actor_critic.action_std
|
||||
entropy_batch = self.actor_critic.entropy
|
||||
|
||||
# KL
|
||||
if self.desired_kl != None and self.schedule == 'adaptive':
|
||||
with torch.inference_mode():
|
||||
kl = torch.sum(
|
||||
torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
|
||||
kl_mean = torch.mean(kl)
|
||||
|
||||
if kl_mean > self.desired_kl * 2.0:
|
||||
self.learning_rate = max(1e-5, self.learning_rate / 1.5)
|
||||
elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
|
||||
self.learning_rate = min(1e-2, self.learning_rate * 1.5)
|
||||
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
|
||||
#Beta VAE loss
|
||||
code,code_vel,decode,mean_vel,logvar_vel,mean_latent,logvar_latent = self.actor_critic.cenet_forward(obs_hist_batch)
|
||||
|
||||
vel_target = prev_critic_obs_batch[:,45:48]
|
||||
decode_target = obs_batch
|
||||
vel_target.requires_grad = False
|
||||
decode_target.requires_grad = False
|
||||
autoenc_loss = (nn.MSELoss()(code_vel,vel_target) + nn.MSELoss()(decode,decode_target) + beta*(-0.5 * torch.sum(1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp())))/self.num_mini_batches
|
||||
# estimation_loss = (code[:,0:3] - prev_critic_obs_batch[:,45:48]).pow(2).mean()
|
||||
# reconst_loss = (decode - obs_batch).pow(2).mean()
|
||||
# latent_loss = beta*(-0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp()))/mean.shape[0]
|
||||
# autoenc_loss = estimation_loss + reconst_loss + latent_loss
|
||||
# Surrogate loss
|
||||
ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
|
||||
surrogate = -torch.squeeze(advantages_batch) * ratio
|
||||
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
|
||||
1.0 + self.clip_param)
|
||||
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
|
||||
|
||||
# Value function loss
|
||||
if self.use_clipped_value_loss:
|
||||
value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
|
||||
self.clip_param)
|
||||
value_losses = (value_batch - returns_batch).pow(2)
|
||||
value_losses_clipped = (value_clipped - returns_batch).pow(2)
|
||||
value_loss = torch.max(value_losses, value_losses_clipped).mean()
|
||||
else:
|
||||
value_loss = (returns_batch - value_batch).pow(2).mean()
|
||||
|
||||
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + autoenc_loss
|
||||
|
||||
# Gradient step
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_autoenc_loss += autoenc_loss.item()
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, mean_autoenc_loss
|
||||
246
motrix_rl/src/motrix_rl/dwaq_rsl/rollout_storage.py
Normal file
246
motrix_rl/src/motrix_rl/dwaq_rsl/rollout_storage.py
Normal file
@@ -0,0 +1,246 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from .utils import split_and_pad_trajectories
|
||||
|
||||
class RolloutStorage:
|
||||
class Transition:
|
||||
def __init__(self):
|
||||
self.observations = None
|
||||
self.critic_observations = None
|
||||
self.prev_critic_obs = None
|
||||
self.observation_history = None
|
||||
self.actions = None
|
||||
self.rewards = None
|
||||
self.dones = None
|
||||
self.values = None
|
||||
self.actions_log_prob = None
|
||||
self.action_mean = None
|
||||
self.action_sigma = None
|
||||
self.hidden_states = None
|
||||
|
||||
def clear(self):
|
||||
self.__init__()
|
||||
|
||||
def __init__(self, num_envs, num_transitions_per_env, obs_shape, privileged_obs_shape, obs_hist_shape, actions_shape, device="cpu"):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.obs_shape = obs_shape
|
||||
self.privileged_obs_shape = privileged_obs_shape
|
||||
self.actions_shape = actions_shape
|
||||
|
||||
# Core
|
||||
self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device)
|
||||
if privileged_obs_shape[0] is not None:
|
||||
self.privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
else:
|
||||
self.privileged_observations = None
|
||||
|
||||
self.prev_critic_obs = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
self.observation_history = torch.zeros(num_transitions_per_env, num_envs, *obs_hist_shape, device=self.device)
|
||||
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
|
||||
|
||||
# For PPO
|
||||
self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
|
||||
self.num_transitions_per_env = num_transitions_per_env
|
||||
self.num_envs = num_envs
|
||||
|
||||
# rnn
|
||||
self.saved_hidden_states_a = None
|
||||
self.saved_hidden_states_c = None
|
||||
|
||||
self.step = 0
|
||||
|
||||
def add_transitions(self, transition: Transition):
|
||||
if self.step >= self.num_transitions_per_env:
|
||||
raise AssertionError("Rollout buffer overflow")
|
||||
self.observations[self.step].copy_(transition.observations)
|
||||
if self.privileged_observations is not None: self.privileged_observations[self.step].copy_(transition.critic_observations)
|
||||
self.prev_critic_obs[self.step].copy_(transition.prev_critic_obs)
|
||||
self.observation_history[self.step].copy_(transition.observation_history)
|
||||
self.actions[self.step].copy_(transition.actions)
|
||||
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
|
||||
self.dones[self.step].copy_(transition.dones.view(-1, 1))
|
||||
self.values[self.step].copy_(transition.values)
|
||||
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
|
||||
self.mu[self.step].copy_(transition.action_mean)
|
||||
self.sigma[self.step].copy_(transition.action_sigma)
|
||||
self._save_hidden_states(transition.hidden_states)
|
||||
self.step += 1
|
||||
|
||||
def _save_hidden_states(self, hidden_states):
|
||||
if hidden_states is None or hidden_states==(None, None):
|
||||
return
|
||||
# make a tuple out of GRU hidden state sto match the LSTM format
|
||||
hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
|
||||
hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
|
||||
|
||||
# initialize if needed
|
||||
if self.saved_hidden_states_a is None:
|
||||
self.saved_hidden_states_a = [torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a))]
|
||||
self.saved_hidden_states_c = [torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c))]
|
||||
# copy the states
|
||||
for i in range(len(hid_a)):
|
||||
self.saved_hidden_states_a[i][self.step].copy_(hid_a[i])
|
||||
self.saved_hidden_states_c[i][self.step].copy_(hid_c[i])
|
||||
|
||||
|
||||
def clear(self):
|
||||
self.step = 0
|
||||
|
||||
def compute_returns(self, last_values, gamma, lam):
|
||||
advantage = 0
|
||||
for step in reversed(range(self.num_transitions_per_env)):
|
||||
if step == self.num_transitions_per_env - 1:
|
||||
next_values = last_values
|
||||
else:
|
||||
next_values = self.values[step + 1]
|
||||
next_is_not_terminal = 1.0 - self.dones[step].float()
|
||||
delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
|
||||
advantage = delta + next_is_not_terminal * gamma * lam * advantage
|
||||
self.returns[step] = advantage + self.values[step]
|
||||
|
||||
# Compute and normalize the advantages
|
||||
self.advantages = self.returns - self.values
|
||||
self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
|
||||
|
||||
def get_statistics(self):
|
||||
done = self.dones
|
||||
done[-1] = 1
|
||||
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]))
|
||||
trajectory_lengths = (done_indices[1:] - done_indices[:-1])
|
||||
return trajectory_lengths.float().mean(), self.rewards.mean()
|
||||
|
||||
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
batch_size = self.num_envs * self.num_transitions_per_env
|
||||
mini_batch_size = batch_size // num_mini_batches
|
||||
indices = torch.randperm(num_mini_batches*mini_batch_size, requires_grad=False, device=self.device)
|
||||
|
||||
observations = self.observations.flatten(0, 1)
|
||||
if self.privileged_observations is not None:
|
||||
critic_observations = self.privileged_observations.flatten(0, 1)
|
||||
else:
|
||||
critic_observations = observations
|
||||
|
||||
prev_critic_obs = self.prev_critic_obs.flatten(0,1)
|
||||
obs_history = self.observation_history.flatten(0,1)
|
||||
actions = self.actions.flatten(0, 1)
|
||||
values = self.values.flatten(0, 1)
|
||||
returns = self.returns.flatten(0, 1)
|
||||
old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
|
||||
advantages = self.advantages.flatten(0, 1)
|
||||
old_mu = self.mu.flatten(0, 1)
|
||||
old_sigma = self.sigma.flatten(0, 1)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
|
||||
start = i*mini_batch_size
|
||||
end = (i+1)*mini_batch_size
|
||||
batch_idx = indices[start:end]
|
||||
|
||||
obs_batch = observations[batch_idx]
|
||||
critic_observations_batch = critic_observations[batch_idx]
|
||||
prev_critic_obs_batch = prev_critic_obs[batch_idx]
|
||||
obs_hist_batch = obs_history[batch_idx]
|
||||
actions_batch = actions[batch_idx]
|
||||
target_values_batch = values[batch_idx]
|
||||
returns_batch = returns[batch_idx]
|
||||
old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
|
||||
advantages_batch = advantages[batch_idx]
|
||||
old_mu_batch = old_mu[batch_idx]
|
||||
old_sigma_batch = old_sigma[batch_idx]
|
||||
yield obs_batch, critic_observations_batch,prev_critic_obs_batch,obs_hist_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (None, None), None
|
||||
|
||||
# for RNNs only
|
||||
def reccurent_mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
|
||||
padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones)
|
||||
if self.privileged_observations is not None:
|
||||
padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones)
|
||||
else:
|
||||
padded_critic_obs_trajectories = padded_obs_trajectories
|
||||
|
||||
mini_batch_size = self.num_envs // num_mini_batches
|
||||
for ep in range(num_epochs):
|
||||
first_traj = 0
|
||||
for i in range(num_mini_batches):
|
||||
start = i*mini_batch_size
|
||||
stop = (i+1)*mini_batch_size
|
||||
|
||||
dones = self.dones.squeeze(-1)
|
||||
last_was_done = torch.zeros_like(dones, dtype=torch.bool)
|
||||
last_was_done[1:] = dones[:-1]
|
||||
last_was_done[0] = True
|
||||
trajectories_batch_size = torch.sum(last_was_done[:, start:stop])
|
||||
last_traj = first_traj + trajectories_batch_size
|
||||
|
||||
masks_batch = trajectory_masks[:, first_traj:last_traj]
|
||||
obs_batch = padded_obs_trajectories[:, first_traj:last_traj]
|
||||
critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj]
|
||||
|
||||
actions_batch = self.actions[:, start:stop]
|
||||
old_mu_batch = self.mu[:, start:stop]
|
||||
old_sigma_batch = self.sigma[:, start:stop]
|
||||
returns_batch = self.returns[:, start:stop]
|
||||
advantages_batch = self.advantages[:, start:stop]
|
||||
values_batch = self.values[:, start:stop]
|
||||
old_actions_log_prob_batch = self.actions_log_prob[:, start:stop]
|
||||
|
||||
# reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim])
|
||||
# then take only time steps after dones (flattens num envs and time dimensions),
|
||||
# take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim]
|
||||
last_was_done = last_was_done.permute(1, 0)
|
||||
hid_a_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
|
||||
for saved_hidden_states in self.saved_hidden_states_a ]
|
||||
hid_c_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
|
||||
for saved_hidden_states in self.saved_hidden_states_c ]
|
||||
# remove the tuple for GRU
|
||||
hid_a_batch = hid_a_batch[0] if len(hid_a_batch)==1 else hid_a_batch
|
||||
hid_c_batch = hid_c_batch[0] if len(hid_c_batch)==1 else hid_a_batch
|
||||
|
||||
yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (hid_a_batch, hid_c_batch), masks_batch
|
||||
|
||||
first_traj = last_traj
|
||||
180
motrix_rl/src/motrix_rl/dwaq_rsl/runner.py
Normal file
180
motrix_rl/src/motrix_rl/dwaq_rsl/runner.py
Normal file
@@ -0,0 +1,180 @@
|
||||
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
||||
"""DreamWaQ OnPolicyRunner — adapted from upstream rsl_rl-1.0.2.
|
||||
|
||||
Faithful port of Manaro-Alpha/DreamWaQ's on_policy_runner.py, with the
|
||||
rsl_rl-package imports replaced by local (relative) imports and the isaacgym
|
||||
noise removed. Drives the PPO + CENet (VAE) joint training loop.
|
||||
"""
|
||||
import time
|
||||
import os
|
||||
from collections import deque
|
||||
import statistics
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
import torch
|
||||
|
||||
from .ppo import PPO
|
||||
from .actor_critic_dwaq import ActorCritic_DWAQ
|
||||
|
||||
|
||||
class OnPolicyRunner:
|
||||
def __init__(self, env, train_cfg, log_dir=None, device="cpu"):
|
||||
self.cfg = train_cfg["runner"]
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
|
||||
num_critic_obs = env.num_privileged_obs if env.num_privileged_obs is not None else env.num_obs
|
||||
cenet_in_dim = env.num_obs_hist * env.num_obs # 5*45 = 225
|
||||
cenet_out_dim = 19 # vel(3) + latent(16)
|
||||
|
||||
actor_critic = ActorCritic_DWAQ(
|
||||
env.num_obs + cenet_out_dim, # actor input: obs(45) + code(19) = 64
|
||||
num_critic_obs, # critic input: privileged_obs
|
||||
env.num_actions, # 12
|
||||
cenet_in_dim, # 225
|
||||
cenet_out_dim, # 19
|
||||
**self.policy_cfg,
|
||||
).to(self.device)
|
||||
|
||||
self.alg = PPO(actor_critic, device=self.device, **self.alg_cfg)
|
||||
self.num_steps_per_env = self.cfg["num_steps_per_env"]
|
||||
self.save_interval = self.cfg["save_interval"]
|
||||
|
||||
self.alg.init_storage(
|
||||
env.num_envs, self.num_steps_per_env,
|
||||
[env.num_obs], [env.num_privileged_obs],
|
||||
[env.num_obs_hist * env.num_obs], [env.num_actions],
|
||||
)
|
||||
|
||||
self.log_dir = log_dir
|
||||
self.writer = None
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.current_learning_iteration = 0
|
||||
|
||||
_, _, _, _ = self.env.reset()
|
||||
|
||||
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
|
||||
if self.log_dir is not None and self.writer is None:
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
|
||||
if init_at_random_ep_len:
|
||||
self.env.episode_length_buf = torch.randint_like(
|
||||
self.env.episode_length_buf, high=int(self.env.max_episode_length))
|
||||
|
||||
obs, obs_hist = self.env.get_observations()
|
||||
privileged_obs, prev_critic_obs = self.env.get_privileged_observations()
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs, prev_critic_obs, obs_hist = (
|
||||
obs.to(self.device), critic_obs.to(self.device),
|
||||
prev_critic_obs.to(self.device), obs_hist.to(self.device))
|
||||
self.alg.actor_critic.train()
|
||||
|
||||
ep_infos = []
|
||||
rewbuffer = deque(maxlen=100)
|
||||
lenbuffer = deque(maxlen=100)
|
||||
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
tot_iter = self.current_learning_iteration + num_learning_iterations
|
||||
for it in range(self.current_learning_iteration, tot_iter):
|
||||
start = time.time()
|
||||
with torch.inference_mode():
|
||||
for i in range(self.num_steps_per_env):
|
||||
actions = self.alg.act(obs, critic_obs, prev_critic_obs, obs_hist)
|
||||
obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, infos = self.env.step(actions)
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
prev_critic_obs = prev_privileged_obs
|
||||
obs, critic_obs, prev_critic_obs, obs_hist, rewards, dones = (
|
||||
obs.to(self.device), critic_obs.to(self.device),
|
||||
prev_critic_obs.to(self.device), obs_hist.to(self.device),
|
||||
rewards.to(self.device), dones.to(self.device))
|
||||
self.alg.process_env_step(rewards, dones, infos)
|
||||
|
||||
if self.log_dir is not None:
|
||||
if "episode" in infos:
|
||||
ep_infos.append(infos["episode"])
|
||||
cur_reward_sum += rewards
|
||||
cur_episode_length += 1
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
|
||||
cur_reward_sum[new_ids] = 0
|
||||
cur_episode_length[new_ids] = 0
|
||||
|
||||
stop = time.time()
|
||||
collection_time = stop - start
|
||||
start = stop
|
||||
self.alg.compute_returns(critic_obs)
|
||||
|
||||
mean_value_loss, mean_surrogate_loss, mean_autoenc_loss = self.alg.update()
|
||||
stop = time.time()
|
||||
learn_time = stop - start
|
||||
if self.log_dir is not None:
|
||||
self.log(locals())
|
||||
if it % self.save_interval == 0:
|
||||
self.save(os.path.join(self.log_dir, "model_{}.pt".format(it)))
|
||||
ep_infos.clear()
|
||||
|
||||
self.current_learning_iteration += num_learning_iterations
|
||||
self.save(os.path.join(self.log_dir, "model_{}.pt".format(self.current_learning_iteration)))
|
||||
|
||||
def log(self, locs, width=80, pad=35):
|
||||
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
|
||||
self.tot_time += locs["collection_time"] + locs["learn_time"]
|
||||
iteration_time = locs["collection_time"] + locs["learn_time"]
|
||||
|
||||
if locs["ep_infos"]:
|
||||
for key in locs["ep_infos"][0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
for ep_info in locs["ep_infos"]:
|
||||
v = ep_info[key]
|
||||
if not isinstance(v, torch.Tensor):
|
||||
v = torch.Tensor([v])
|
||||
if len(v.shape) == 0:
|
||||
v = v.unsqueeze(0)
|
||||
infotensor = torch.cat((infotensor, v.to(self.device)))
|
||||
self.writer.add_scalar("Episode/" + key, torch.mean(infotensor), locs["it"])
|
||||
mean_std = self.alg.actor_critic.std.mean()
|
||||
fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"]))
|
||||
|
||||
self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"])
|
||||
self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"])
|
||||
self.writer.add_scalar("Loss/autoenc_function", locs["mean_autoenc_loss"], locs["it"])
|
||||
self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"])
|
||||
self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"])
|
||||
self.writer.add_scalar("Perf/total_fps", fps, locs["it"])
|
||||
if len(locs["rewbuffer"]) > 0:
|
||||
self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"])
|
||||
self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"])
|
||||
|
||||
if it_is_log := (locs["it"] % 10 == 0):
|
||||
rew = statistics.mean(locs["rewbuffer"]) if len(locs["rewbuffer"]) > 0 else 0.0
|
||||
elen = statistics.mean(locs["lenbuffer"]) if len(locs["lenbuffer"]) > 0 else 0.0
|
||||
print(f"[it {locs['it']}/{self.current_learning_iteration + locs['num_learning_iterations']}] "
|
||||
f"fps={fps} val_loss={locs['mean_value_loss']:.3f} surr={locs['mean_surrogate_loss']:.4f} "
|
||||
f"vae={locs['mean_autoenc_loss']:.3f} std={mean_std.item():.2f} "
|
||||
f"rew={rew:.2f} eplen={elen:.0f}")
|
||||
|
||||
def save(self, path, infos=None):
|
||||
torch.save({
|
||||
"model_state_dict": self.alg.actor_critic.state_dict(),
|
||||
"optimizer_state_dict": self.alg.optimizer.state_dict(),
|
||||
"iter": self.current_learning_iteration,
|
||||
"infos": infos,
|
||||
}, path)
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
loaded_dict = torch.load(path, map_location=self.device)
|
||||
self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"])
|
||||
if load_optimizer:
|
||||
self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
|
||||
self.current_learning_iteration = loaded_dict["iter"]
|
||||
return loaded_dict["infos"]
|
||||
|
||||
def get_inference_policy(self, device=None):
|
||||
self.alg.actor_critic.eval()
|
||||
if device is not None:
|
||||
self.alg.actor_critic.to(device)
|
||||
return self.alg.actor_critic.act_inference
|
||||
71
motrix_rl/src/motrix_rl/dwaq_rsl/utils.py
Normal file
71
motrix_rl/src/motrix_rl/dwaq_rsl/utils.py
Normal file
@@ -0,0 +1,71 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
|
||||
def split_and_pad_trajectories(tensor, dones):
|
||||
""" Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory.
|
||||
Returns masks corresponding to valid parts of the trajectories
|
||||
Example:
|
||||
Input: [ [a1, a2, a3, a4 | a5, a6],
|
||||
[b1, b2 | b3, b4, b5 | b6]
|
||||
]
|
||||
|
||||
Output:[ [a1, a2, a3, a4], | [ [True, True, True, True],
|
||||
[a5, a6, 0, 0], | [True, True, False, False],
|
||||
[b1, b2, 0, 0], | [True, True, False, False],
|
||||
[b3, b4, b5, 0], | [True, True, True, False],
|
||||
[b6, 0, 0, 0] | [True, False, False, False],
|
||||
] | ]
|
||||
|
||||
Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions]
|
||||
"""
|
||||
dones = dones.clone()
|
||||
dones[-1] = 1
|
||||
# Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping
|
||||
flat_dones = dones.transpose(1, 0).reshape(-1, 1)
|
||||
|
||||
# Get length of trajectory by counting the number of successive not done elements
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
|
||||
trajectory_lengths = done_indices[1:] - done_indices[:-1]
|
||||
trajectory_lengths_list = trajectory_lengths.tolist()
|
||||
# Extract the individual trajectories
|
||||
trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list)
|
||||
padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)
|
||||
|
||||
|
||||
trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1)
|
||||
return padded_trajectories, trajectory_masks
|
||||
|
||||
def unpad_trajectories(trajectories, masks):
|
||||
""" Does the inverse operation of split_and_pad_trajectories()
|
||||
"""
|
||||
# Need to transpose before and after the masking to have proper reshaping
|
||||
return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0)
|
||||
@@ -116,15 +116,15 @@ class CENetActorModel(MLPModel):
|
||||
self._last_cenet_output = None
|
||||
|
||||
def _update_distribution(self, obs: torch.Tensor) -> None:
|
||||
"""覆盖父类 — 确保 std 始终为正，防止 NaN。"""
|
||||
super()._update_distribution(obs)
|
||||
# 如果 std 因数值问题变负，clamp 到最小值
|
||||
"""覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。"""
|
||||
# 先 clamp std，再调父类创建分布
|
||||
if self.stochastic and not self.state_dependent_std:
|
||||
with torch.no_grad():
|
||||
if self.noise_std_type == "scalar":
|
||||
self.std.clamp_(min=1e-6)
|
||||
elif self.noise_std_type == "log":
|
||||
self.log_std.clamp_(min=-20.0, max=10.0)
|
||||
super()._update_distribution(obs)
|
||||
|
||||
def _get_latent_dim(self) -> int:
|
||||
"""Actor 实际输入：code(19) + policy(45) = 64。"""
|
||||
|
||||
@@ -81,6 +81,9 @@ class DreamWaQPPO(PPO):
|
||||
returns_batch, old_actions_log_prob_batch, old_mu_batch,
|
||||
old_sigma_batch, hid_states_batch, masks_batch,
|
||||
) in generator:
|
||||
# NaN 检测 — 数据有 NaN 就跳过这个 batch
|
||||
if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
|
||||
continue
|
||||
# ── 标准 PPO 前向 ──
|
||||
self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
|
||||
self.critic(obs_batch, masks=masks_batch)
|
||||
|
||||
306
motrix_rl/src/motrix_rl/skrl/jax/train/cts_ppo.py
Normal file
306
motrix_rl/src/motrix_rl/skrl/jax/train/cts_ppo.py
Normal file
@@ -0,0 +1,306 @@
|
||||
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
||||
"""CTS PPO v2 — Teacher/Student encoders with distillation loss.
|
||||
|
||||
Matches go2_rl_gym CTS training flow:
|
||||
Phase 1: Standard PPO update with teacher-computed latent
|
||||
Phase 2: Distillation: MSE(teacher_latent, student_latent) on student envs
|
||||
|
||||
Architecture:
|
||||
Teacher: privileged_obs(76) → MLP(512,256) → L2Norm → latent(32)
|
||||
Student: history_obs(225) → MLP(512,256,256) → L2Norm → latent(32)
|
||||
Actor: student_obs(45) + latent(32) = 77 → [512,256,128] → action(12)
|
||||
Critic: student_obs(45) + latent(32) = 77 → [512,256,128] → value(1)
|
||||
|
||||
The wrapper computes latent online (teacher for 75% envs, student for 25%)
|
||||
and augments observations to 77-dim for standard PPO. After each PPO update,
|
||||
a distillation step trains the student encoder to match the teacher's latent.
|
||||
"""
|
||||
|
||||
import flax.linen as nn
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import optax
|
||||
|
||||
from skrl.agents.jax.ppo import PPO as BasePPO
|
||||
from skrl.memories.jax import RandomMemory
|
||||
from skrl.trainers.jax import SequentialTrainer
|
||||
from skrl.utils import set_seed
|
||||
|
||||
from motrix_envs import registry as env_registry
|
||||
from motrix_rl import registry, utils
|
||||
from motrix_rl.skrl import get_log_dir
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
# Encoder Modules (Flax)
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TeacherEncoder(nn.Module):
|
||||
latent_dim: int = 32
|
||||
hidden_dims: tuple = (512, 256)
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
for d in self.hidden_dims:
|
||||
x = nn.elu(nn.Dense(d)(x))
|
||||
return nn.Dense(self.latent_dim)(x)
|
||||
|
||||
|
||||
class StudentEncoder(nn.Module):
|
||||
latent_dim: int = 32
|
||||
hidden_dims: tuple = (512, 256, 256)
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
for d in self.hidden_dims:
|
||||
x = nn.elu(nn.Dense(d)(x))
|
||||
return nn.Dense(self.latent_dim)(x)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
# CTS Env Wrapper
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CTSEnvWrapper:
|
||||
"""Wraps CTS env. Computes latent online, returns augmented obs (77-dim)."""
|
||||
|
||||
def __init__(self, env, teacher_ratio=0.75, latent_dim=32,
|
||||
teacher_enc=None, teacher_params=None,
|
||||
student_enc=None, student_params=None):
|
||||
self.__env = env
|
||||
self._teacher_ratio = teacher_ratio
|
||||
self._latent_dim = latent_dim
|
||||
self._num_envs = env.num_envs
|
||||
self._t_enc = teacher_enc
|
||||
self._t_params = teacher_params
|
||||
self._s_enc = student_enc
|
||||
self._s_params = student_params
|
||||
from gymnasium import spaces
|
||||
self._obs_space = spaces.Box(-np.inf, np.inf,
|
||||
(45 + latent_dim,), dtype=np.float32)
|
||||
|
||||
@property
|
||||
def num_envs(self): return self._num_envs
|
||||
@property
|
||||
def num_teacher(self): return int(self._num_envs * self._teacher_ratio)
|
||||
@property
|
||||
def teacher_mask(self):
|
||||
m = np.zeros(self._num_envs, dtype=np.bool)
|
||||
m[:self.num_teacher] = True
|
||||
return m
|
||||
@property
|
||||
def student_mask(self): return ~self.teacher_mask
|
||||
@property
|
||||
def privileged_obs(self):
|
||||
return self.__env._state.info.get("privileged_obs",
|
||||
np.zeros((self._num_envs, 76), dtype=np.float32))
|
||||
@property
|
||||
def history_obs(self): return self.__env.history_obs
|
||||
@property
|
||||
def observation_space(self): return self._obs_space
|
||||
@property
|
||||
def action_space(self): return self.__env.action_space
|
||||
@property
|
||||
def device(self): return None
|
||||
@property
|
||||
def num_agents(self): return 1
|
||||
@property
|
||||
def state_preprocessor(self): return None
|
||||
@property
|
||||
def value_preprocessor(self): return None
|
||||
@property
|
||||
def world_info_idx(self): return None
|
||||
@property
|
||||
def random_timestep(self): return False
|
||||
@property
|
||||
def running(self): return True
|
||||
@property
|
||||
def _env(self): return self.__env
|
||||
|
||||
def _augment(self, states):
|
||||
n = states.shape[0]
|
||||
latent = np.zeros((n, self._latent_dim), dtype=np.float32)
|
||||
t_mask = self.teacher_mask
|
||||
if t_mask.any() and self._t_params is not None:
|
||||
priv = jnp.array(self.privileged_obs[t_mask])
|
||||
latent[t_mask] = np.array(self._t_enc.apply(self._t_params, priv))
|
||||
s_mask = self.student_mask
|
||||
if s_mask.any() and self._s_params is not None:
|
||||
hist = jnp.array(self.history_obs[s_mask])
|
||||
latent[s_mask] = np.array(self._s_enc.apply(self._s_params, hist))
|
||||
return np.concatenate([states, latent], axis=-1)
|
||||
|
||||
def reset(self):
|
||||
state = self.__env.init_state()
|
||||
return self._augment(state.obs), state.info
|
||||
|
||||
def step(self, actions):
|
||||
state = self.__env.step(np.array(actions))
|
||||
return (self._augment(state.obs), state.reward.reshape(-1, 1),
|
||||
state.terminated.reshape(-1, 1),
|
||||
state.truncated.reshape(-1, 1), state.info)
|
||||
|
||||
def render(self, *a, **kw): pass
|
||||
def close(self): pass
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
# CTS Agent — adds distillation step after PPO update
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CTSAgent(BasePPO):
|
||||
"""PPO agent with post-update distillation on student encoder."""
|
||||
|
||||
def __init__(self, models, memory, cfg, observation_space, action_space, device,
|
||||
teacher_enc=None, student_enc=None, cts_wrapper=None,
|
||||
teacher_params=None, student_params=None,
|
||||
distill_coef=0.1, student_lr=1e-3):
|
||||
super().__init__(models=models, memory=memory, cfg=cfg,
|
||||
observation_space=observation_space,
|
||||
action_space=action_space, device=device)
|
||||
self._t_enc = teacher_enc
|
||||
self._s_enc = student_enc
|
||||
self._t_params = teacher_params
|
||||
self._s_params = student_params
|
||||
self._wrapper = cts_wrapper
|
||||
self._distill_coef = distill_coef
|
||||
self._latent_dim = getattr(teacher_enc, 'latent_dim', 32) if teacher_enc else 32
|
||||
|
||||
# Separate optimizer for student encoder
|
||||
if student_params is not None:
|
||||
self._s_optimizer = optax.adam(student_lr)
|
||||
self._s_opt_state = self._s_optimizer.init(student_params)
|
||||
else:
|
||||
self._s_optimizer = None
|
||||
|
||||
def update(self):
|
||||
"""PPO update + distillation phase."""
|
||||
result = super().update()
|
||||
|
||||
# Phase 2: Distillation — MSE(teacher_latent, student_latent) on student envs
|
||||
if self._s_params is not None and self._s_optimizer is not None:
|
||||
self._distillation_step()
|
||||
|
||||
return result
|
||||
|
||||
def _distillation_step(self):
|
||||
"""Train student encoder to match teacher latent.
|
||||
|
||||
go2_rl_gym Phase 2: only on student env samples.
|
||||
Loss = || teacher_enc(privileged_obs) - student_enc(history_obs) ||²
|
||||
"""
|
||||
wrapper = self._wrapper
|
||||
if wrapper is None:
|
||||
return
|
||||
|
||||
s_mask = wrapper.student_mask
|
||||
if not s_mask.any():
|
||||
return
|
||||
|
||||
# Get current student env observations
|
||||
priv = jnp.array(wrapper.privileged_obs[s_mask]) # [S, 76]
|
||||
hist = jnp.array(wrapper.history_obs[s_mask]) # [S, 225]
|
||||
|
||||
# L2 normalize both before computing MSE (match go2_rl_gym)
|
||||
def l2_norm(x):
|
||||
return x / (jnp.linalg.norm(x, axis=-1, keepdims=True) + 1e-8)
|
||||
|
||||
def distill_loss(s_params):
|
||||
t_latent = l2_norm(jax.lax.stop_gradient(
|
||||
self._t_enc.apply(self._t_params, priv)))
|
||||
s_latent = l2_norm(self._s_enc.apply(s_params, hist))
|
||||
return jnp.mean(jnp.square(t_latent - s_latent))
|
||||
|
||||
loss, grads = jax.value_and_grad(distill_loss)(self._s_params)
|
||||
updates, self._s_opt_state = self._s_optimizer.update(
|
||||
grads, self._s_opt_state, self._s_params)
|
||||
self._s_params = optax.apply_updates(self._s_params, updates)
|
||||
|
||||
# Track
|
||||
self.tracking_data["Loss / Distillation"].append(loss)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
# Trainer
|
||||
# ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CTSTrainer:
|
||||
def __init__(self, env_name="go1-cts-flat-walk-go2style",
|
||||
sim_backend=None, enable_render=False, cfg_override=None):
|
||||
self._env_name = env_name
|
||||
self._sim_backend = sim_backend
|
||||
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
|
||||
if cfg_override:
|
||||
rlcfg = utils.cfg_override(rlcfg, cfg_override)
|
||||
self._rlcfg = rlcfg
|
||||
|
||||
def train(self):
|
||||
cfg = self._rlcfg
|
||||
set_seed(cfg.runner.seed)
|
||||
|
||||
env = env_registry.make(self._env_name, num_envs=cfg.num_envs)
|
||||
env_cfg = env._cfg
|
||||
latent_dim = getattr(env_cfg, 'latent_dim', 32)
|
||||
priv_dim = getattr(env_cfg, 'num_privileged_obs', 76)
|
||||
hist_len = getattr(env_cfg, 'history_length', 5)
|
||||
student_dim = getattr(env_cfg, 'num_student_obs', 45)
|
||||
teacher_ratio = getattr(env_cfg, 'teacher_env_ratio', 0.75)
|
||||
|
||||
# Init Flax encoders
|
||||
rng = jax.random.PRNGKey(cfg.runner.seed)
|
||||
rng, t_key, s_key = jax.random.split(rng, 3)
|
||||
teacher_enc = TeacherEncoder(latent_dim=latent_dim)
|
||||
student_enc = StudentEncoder(latent_dim=latent_dim)
|
||||
t_params = teacher_enc.init(t_key, jnp.ones((1, priv_dim)))
|
||||
s_params = student_enc.init(s_key, jnp.ones((1, hist_len * student_dim)))
|
||||
|
||||
# CTS wrapper
|
||||
wrapper = CTSEnvWrapper(env, teacher_ratio=teacher_ratio, latent_dim=latent_dim,
|
||||
teacher_enc=teacher_enc, teacher_params=t_params,
|
||||
student_enc=student_enc, student_params=s_params)
|
||||
env.init_state()
|
||||
|
||||
from gymnasium import spaces
|
||||
aug_obs = spaces.Box(-np.inf, np.inf, (student_dim + latent_dim,), dtype=np.float32)
|
||||
wrapper._obs_space = aug_obs
|
||||
|
||||
nT = int(cfg.num_envs * teacher_ratio)
|
||||
nS = cfg.num_envs - nT
|
||||
print(f"[CTS v2] {cfg.num_envs} envs ({nT}T/{nS}S), latent={latent_dim}, "
|
||||
f"priv={priv_dim}, hist={hist_len}×{student_dim}, distill_coef=0.1")
|
||||
|
||||
# Standard model creation
|
||||
from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config
|
||||
std = StdTrainer(self._env_name, sim_backend=None, enable_render=False)
|
||||
std._rlcfg = cfg
|
||||
models = std._make_model(wrapper, cfg)
|
||||
|
||||
ppo_cfg = cfg.runner.agent.to_dict()
|
||||
_add_runtime_config(ppo_cfg, wrapper,
|
||||
log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO"))
|
||||
|
||||
mem_size = cfg.runner.memory.memory_size
|
||||
if mem_size == -1:
|
||||
mem_size = ppo_cfg.get("rollouts", 24)
|
||||
memory = RandomMemory(memory_size=mem_size, num_envs=cfg.num_envs, device=None)
|
||||
|
||||
agent = CTSAgent(
|
||||
models=models, memory=memory, cfg=ppo_cfg,
|
||||
observation_space=aug_obs, action_space=env.action_space,
|
||||
device=None,
|
||||
teacher_enc=teacher_enc, student_enc=student_enc,
|
||||
cts_wrapper=wrapper,
|
||||
teacher_params=t_params, student_params=s_params,
|
||||
distill_coef=0.1, student_lr=1e-3,
|
||||
)
|
||||
|
||||
trainer_cfg = {"timesteps": cfg.runner.trainer.timesteps, "headless": True}
|
||||
SequentialTrainer(cfg=trainer_cfg, env=wrapper, agents=agent).train()
|
||||
|
||||
# Save encoder params for deployment
|
||||
import pickle as _pk
|
||||
_log = get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO")
|
||||
_pk.dump(t_params, open(f"{_log}/teacher_params.pkl", "wb"))
|
||||
_pk.dump(s_params, open(f"{_log}/student_params.pkl", "wb"))
|
||||
print(f"[CTS] Encoder weights saved → {_log}")
|
||||
540
motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py
Normal file
540
motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py
Normal file
@@ -0,0 +1,540 @@
|
||||
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
|
||||
"""DreamWaQ PPO — aligned with Manaro-Alpha/DreamWaQ.
|
||||
|
||||
CENet (VAE): history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
|
||||
Decoder: code(19) → [64,128] → next_obs(45)
|
||||
Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12)
|
||||
Critic: privileged_obs(235) → [512,256,128] → value(1)
|
||||
|
||||
VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL
|
||||
"""
|
||||
|
||||
import flax.linen as nn
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import optax
|
||||
import pickle, os
|
||||
|
||||
from skrl.agents.jax.ppo import PPO as BasePPO
|
||||
from skrl.memories.jax import RandomMemory
|
||||
from skrl.models.jax import DeterministicMixin, GaussianMixin, Model
|
||||
from skrl.trainers.jax import SequentialTrainer
|
||||
from skrl.utils import set_seed
|
||||
|
||||
from motrix_envs import registry as env_registry
|
||||
from motrix_envs.locomotion.go1.dreamwaq import DreamWaQTask
|
||||
from motrix_rl import registry, utils
|
||||
from motrix_rl.skrl import get_log_dir
|
||||
|
||||
# Ensure RL config registered
|
||||
import motrix_rl.tasks.go1_dreamwaq # noqa: F401
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# CENet (VAE) — exact Manaro-Alpha architecture
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CENet(nn.Module):
|
||||
"""VAE for terrain imagination + velocity estimation.
|
||||
|
||||
Encoder: 225 → Linear(128) → ELU → Linear(64) → ELU
|
||||
4 heads on 64-dim bottleneck:
|
||||
- latent_mu, latent_logvar: 64 → 16 each
|
||||
- vel_mu, vel_logvar: 64 → 3 each
|
||||
Decoder: (latent+vel=19) → Linear(64) → ELU → Linear(128) → ELU → Linear(45)
|
||||
"""
|
||||
num_obs: int = 45
|
||||
num_history: int = 5
|
||||
num_latent: int = 16 # latent z dim
|
||||
cenet_out_dim: int = 19 # code = vel(3) + latent(16)
|
||||
|
||||
def setup(self):
|
||||
# Orthogonal weight init matching upstream torch.nn.init.orthogonal_
|
||||
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
|
||||
|
||||
# Encoder (matches Manaro-Alpha: 225 → 128 → 64)
|
||||
self.enc_fc1 = nn.Dense(128, kernel_init=kernel_init)
|
||||
self.enc_fc2 = nn.Dense(64, kernel_init=kernel_init)
|
||||
|
||||
# 4 output heads
|
||||
self.latent_mu = nn.Dense(self.num_latent, kernel_init=kernel_init)
|
||||
self.latent_logvar = nn.Dense(self.num_latent, kernel_init=kernel_init)
|
||||
self.vel_mu = nn.Dense(3, kernel_init=kernel_init)
|
||||
self.vel_logvar = nn.Dense(3, kernel_init=kernel_init)
|
||||
|
||||
# Decoder (matches Manaro-Alpha: 19 → 64 → 128 → 45)
|
||||
self.dec_fc1 = nn.Dense(64, kernel_init=kernel_init)
|
||||
self.dec_fc2 = nn.Dense(128, kernel_init=kernel_init)
|
||||
self.dec_out = nn.Dense(self.num_obs, kernel_init=kernel_init)
|
||||
|
||||
def encode(self, obs_history):
|
||||
"""obs_history: (B, T, 45) or (B, T*45) → (latent_mu, latent_logvar, vel_mu, vel_logvar)."""
|
||||
x = obs_history.reshape(obs_history.shape[0], -1)
|
||||
x = nn.elu(self.enc_fc1(x))
|
||||
x = nn.elu(self.enc_fc2(x))
|
||||
return (self.latent_mu(x), self.latent_logvar(x),
|
||||
self.vel_mu(x), self.vel_logvar(x))
|
||||
|
||||
def reparameterize(self, mu, logvar, rng):
|
||||
std = jnp.exp(0.5 * logvar)
|
||||
return mu + std * jax.random.normal(rng, mu.shape)
|
||||
|
||||
def decode(self, z, vel):
|
||||
x = jnp.concatenate([z, vel], axis=-1)
|
||||
x = nn.elu(self.dec_fc1(x))
|
||||
x = nn.elu(self.dec_fc2(x))
|
||||
return self.dec_out(x)
|
||||
|
||||
def __call__(self, obs_history, rng):
|
||||
lm, ll, vm, vl = self.encode(obs_history)
|
||||
rng_z, rng_v = jax.random.split(rng)
|
||||
z = self.reparameterize(lm, ll, rng_z)
|
||||
vel = self.reparameterize(vm, vl, rng_v)
|
||||
return {
|
||||
'z': z, 'vel_est': vel,
|
||||
'latent_mu': lm, 'latent_logvar': ll,
|
||||
'vel_mu': vm, 'vel_logvar': vl,
|
||||
'obs_pred': self.decode(z, vel),
|
||||
}
|
||||
|
||||
def sample(self, obs_history, rng):
|
||||
"""Sample (z, vel) for rollout (adds noise)."""
|
||||
lm, ll, vm, vl = self.encode(obs_history)
|
||||
rng_z, rng_v = jax.random.split(rng)
|
||||
z = self.reparameterize(lm, ll, rng_z)
|
||||
vel = self.reparameterize(vm, vl, rng_v)
|
||||
return z, vel
|
||||
|
||||
def inference(self, obs_history):
|
||||
"""Mean (z, vel) for deployment (no noise)."""
|
||||
lm, ll, vm, vl = self.encode(obs_history)
|
||||
return lm, vm
|
||||
|
||||
def loss_fn(self, obs_history, current_obs, base_vel, rng, kl_weight=1.0):
|
||||
"""VAE loss: reconstructs current obs (last history frame), NOT next obs.
|
||||
|
||||
Upstream: decode_target = obs_batch (the same timestep's observation,
|
||||
not a future prediction). This is an autoencoder over current observation.
|
||||
"""
|
||||
lm, ll, vm, vl = self.encode(obs_history)
|
||||
rng_z, rng_v = jax.random.split(rng)
|
||||
z = self.reparameterize(lm, ll, rng_z)
|
||||
vel = self.reparameterize(vm, vl, rng_v)
|
||||
pred = self.decode(z, vel)
|
||||
|
||||
recons = jnp.mean(jnp.square(pred - current_obs), axis=-1)
|
||||
vel_loss = jnp.mean(jnp.square(vel - base_vel), axis=-1)
|
||||
kld = -0.5 * jnp.sum(1 + ll - jnp.square(lm) - jnp.exp(ll), axis=-1)
|
||||
|
||||
return {
|
||||
'loss': jnp.mean(recons + vel_loss + kl_weight * kld),
|
||||
'recons_loss': jnp.mean(recons),
|
||||
'vel_loss': jnp.mean(vel_loss),
|
||||
'kld_loss': jnp.mean(kld),
|
||||
}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Environment Wrapper
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class DreamWaQWrapper:
|
||||
"""Wraps DreamWaQTask, computes CENet code online.
|
||||
|
||||
Augmented observation layout (254-dim):
|
||||
[code_vel(3) | code_latent(16) | obs(45) | base_vel(3) | heights(187)]
|
||||
|
||||
Actor uses: [0:64] = code(19) + obs(45)
|
||||
Critic uses: [19:254] = obs(45) + base_vel(3) + heights(187)
|
||||
"""
|
||||
|
||||
def __init__(self, env, cenet, cenet_params, rng=None, vae_buf_size=24):
|
||||
self._env = env
|
||||
self._cenet = cenet
|
||||
self._cenet_params = cenet_params
|
||||
self._num_envs = env.num_envs
|
||||
self._rng = rng if rng is not None else jax.random.PRNGKey(0)
|
||||
self._code_dim = 19 # vel(3) + latent(16)
|
||||
self._aug_dim = self._code_dim + 45 + 3 + 187 # 254
|
||||
|
||||
from gymnasium import spaces
|
||||
self._obs_space = spaces.Box(-np.inf, np.inf, (self._aug_dim,), dtype=np.float32)
|
||||
|
||||
# VAE training buffer
|
||||
self._vae_buf = []
|
||||
self._vae_buf_size = vae_buf_size
|
||||
|
||||
@property
|
||||
def num_envs(self): return self._num_envs
|
||||
@property
|
||||
def observation_space(self): return self._obs_space
|
||||
@property
|
||||
def action_space(self): return self._env.action_space
|
||||
@property
|
||||
def device(self): return None
|
||||
@property
|
||||
def num_agents(self): return 1
|
||||
@property
|
||||
def state_preprocessor(self): return None
|
||||
@property
|
||||
def value_preprocessor(self): return None
|
||||
@property
|
||||
def world_info_idx(self): return None
|
||||
@property
|
||||
def random_timestep(self): return False
|
||||
@property
|
||||
def running(self): return True
|
||||
|
||||
@property
|
||||
def privileged_obs(self):
|
||||
return self._env._state.info.get("privileged_obs",
|
||||
np.zeros((self._num_envs, 235), dtype=np.float32))
|
||||
|
||||
@property
|
||||
def obs_history(self):
|
||||
return self._env._state.info.get("obs_history",
|
||||
np.zeros((self._num_envs, 5, 45), dtype=np.float32))
|
||||
|
||||
@property
|
||||
def base_vel(self):
|
||||
return self._env._state.info.get("base_vel",
|
||||
np.zeros((self._num_envs, 3), dtype=np.float32))
|
||||
|
||||
@property
|
||||
def vae_training_data(self):
|
||||
return self._vae_buf
|
||||
|
||||
def _augment(self, obs):
|
||||
"""Compute CENet code and build augmented obs."""
|
||||
hist = jnp.array(self.obs_history)
|
||||
self._rng, rng_key = jax.random.split(self._rng)
|
||||
z, vel = self._cenet.apply(self._cenet_params, hist, rng_key, method=self._cenet.sample)
|
||||
code = np.concatenate([np.array(vel), np.array(z)], axis=-1) # (N, 19)
|
||||
|
||||
priv = self.privileged_obs
|
||||
heights = priv[:, 45+3:] # (N, 187)
|
||||
base_vel_norm = self.base_vel * np.array([2.0, 2.0, 1.0], dtype=np.float32)
|
||||
|
||||
return np.concatenate([code, obs, base_vel_norm, heights], axis=-1)
|
||||
|
||||
def reset(self):
|
||||
state = self._env.init_state()
|
||||
self._vae_buf = []
|
||||
return self._augment(state.obs), state.info
|
||||
|
||||
def step(self, actions):
|
||||
old_hist = self.obs_history.copy()
|
||||
state = self._env.step(np.array(actions, dtype=np.float32))
|
||||
|
||||
# Store for VAE training: history → current_obs (last frame of history)
|
||||
# Upstream: decode target = obs_batch (current observation, not next)
|
||||
self._vae_buf.append((old_hist, old_hist[:, -1, :].copy(), self.base_vel.copy()))
|
||||
if len(self._vae_buf) > self._vae_buf_size:
|
||||
self._vae_buf.pop(0)
|
||||
|
||||
return (self._augment(state.obs),
|
||||
state.reward.reshape(-1, 1).astype(np.float32),
|
||||
state.terminated.reshape(-1, 1).astype(np.float32),
|
||||
state.truncated.reshape(-1, 1).astype(np.float32),
|
||||
state.info)
|
||||
|
||||
def render(self, *a, **kw): pass
|
||||
def close(self): pass
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# DreamWaQ Agent — PPO + CENet training
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class DreamWaQAgent(BasePPO):
|
||||
"""PPO agent with asymmetric critic + CENet (VAE) training."""
|
||||
|
||||
def __init__(self, models, memory, cfg, observation_space, action_space, device,
|
||||
cenet=None, cenet_params=None, wrapper=None,
|
||||
kl_weight=1.0, vae_lr=1e-3, save_dir=None):
|
||||
super().__init__(models=models, memory=memory, cfg=cfg,
|
||||
observation_space=observation_space,
|
||||
action_space=action_space, device=device)
|
||||
self._cenet = cenet
|
||||
self._cenet_params = cenet_params
|
||||
self._wrapper = wrapper
|
||||
self._kl_weight = kl_weight
|
||||
self._save_dir = save_dir
|
||||
self._vae_count = 0
|
||||
self._rng = jax.random.PRNGKey(42)
|
||||
|
||||
if cenet_params is not None:
|
||||
self._vae_opt = optax.adam(vae_lr)
|
||||
self._vae_opt_state = self._vae_opt.init(cenet_params)
|
||||
else:
|
||||
self._vae_opt = None
|
||||
|
||||
# Override SKRL's internal _update (NOT 'update')
|
||||
def _update(self, timestep: int, timesteps: int):
|
||||
super()._update(timestep, timesteps)
|
||||
|
||||
# Track detailed reward terms
|
||||
wrapper = self._wrapper
|
||||
if wrapper is not None:
|
||||
try:
|
||||
terms = wrapper._env._state.info.get("reward_terms", {})
|
||||
for name, val in terms.items():
|
||||
self.tracking_data[f"Reward / {name}"].append(float(val))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# CENet training after each PPO update
|
||||
if self._cenet_params is not None and self._vae_opt is not None:
|
||||
self._vae_step()
|
||||
|
||||
def _vae_step(self):
|
||||
"""Train CENet: reconstruct current obs (last history frame) + velocity supervision + KL."""
|
||||
wrapper = self._wrapper
|
||||
if wrapper is None or len(wrapper.vae_training_data) == 0:
|
||||
return
|
||||
|
||||
hist_np, current_obs_np, base_vel_np = wrapper.vae_training_data[-1]
|
||||
hist = jnp.array(hist_np)
|
||||
current_obs = jnp.array(current_obs_np)
|
||||
base_vel = jnp.array(base_vel_np)
|
||||
|
||||
n = wrapper.num_envs
|
||||
valid = jnp.linalg.norm(hist.reshape(n, -1), axis=-1) > 1e-6
|
||||
if not valid.any():
|
||||
return
|
||||
|
||||
self._rng, vae_rng = jax.random.split(self._rng)
|
||||
|
||||
def loss_fn(p):
|
||||
return self._cenet.apply(
|
||||
p, hist, current_obs, base_vel, vae_rng, self._kl_weight,
|
||||
method=self._cenet.loss_fn)['loss']
|
||||
|
||||
loss, grads = jax.value_and_grad(loss_fn)(self._cenet_params)
|
||||
updates, self._vae_opt_state = self._vae_opt.update(
|
||||
grads, self._vae_opt_state, self._cenet_params)
|
||||
self._cenet_params = optax.apply_updates(self._cenet_params, updates)
|
||||
|
||||
# Sync wrapper
|
||||
wrapper._cenet_params = self._cenet_params
|
||||
|
||||
self._vae_count += 1
|
||||
if self._save_dir and self._vae_count % 200 == 0:
|
||||
pickle.dump(self._cenet_params,
|
||||
open(f"{self._save_dir}/vae_{self._vae_count}.pkl", "wb"))
|
||||
|
||||
self.tracking_data["Loss / VAE"].append(loss)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Trainer
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class DreamWaQTrainer:
|
||||
def __init__(self, env_name="go1-dreamwaq-walk", cfg_override=None):
|
||||
self._env_name = env_name
|
||||
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
|
||||
if cfg_override:
|
||||
rlcfg = utils.cfg_override(rlcfg, cfg_override)
|
||||
self._rlcfg = rlcfg
|
||||
|
||||
def train(self):
|
||||
cfg = self._rlcfg
|
||||
set_seed(cfg.runner.seed)
|
||||
|
||||
env = env_registry.make(self._env_name, num_envs=cfg.num_envs)
|
||||
env_cfg = env._cfg
|
||||
|
||||
# Init CENet
|
||||
rng = jax.random.PRNGKey(cfg.runner.seed)
|
||||
rng, crng = jax.random.split(rng)
|
||||
cenet = CENet(num_obs=45, num_history=env_cfg.num_history,
|
||||
num_latent=env_cfg.num_latent,
|
||||
cenet_out_dim=env_cfg.cenet_out_dim)
|
||||
cenet_params = cenet.init(crng, jnp.ones((1, env_cfg.num_history, 45)), jax.random.PRNGKey(0))
|
||||
|
||||
# Wrapper
|
||||
wrapper = DreamWaQWrapper(env, cenet, cenet_params,
|
||||
rng=jax.random.PRNGKey(cfg.runner.seed + 1))
|
||||
env.init_state()
|
||||
|
||||
print(f"[DreamWaQ] {cfg.num_envs} envs, CENet[128,64], latent=16+3=19, "
|
||||
f"aug_dim={wrapper.observation_space.shape[0]}")
|
||||
|
||||
# Build models
|
||||
from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config
|
||||
|
||||
policy_cfg = cfg.runner.models.policy
|
||||
value_cfg = cfg.runner.models.value
|
||||
|
||||
class Policy(GaussianMixin, Model):
|
||||
def __init__(self, observation_space, action_space, device=None, **kwargs):
|
||||
Model.__init__(self, observation_space, action_space, device, **kwargs)
|
||||
GaussianMixin.__init__(self, policy_cfg.clip_actions,
|
||||
policy_cfg.clip_log_std,
|
||||
policy_cfg.min_log_std, policy_cfg.max_log_std,
|
||||
policy_cfg.reduction)
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, inputs, role):
|
||||
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
|
||||
x = inputs["states"]
|
||||
# Actor: code(19) + obs(45) = first 64 dims
|
||||
x = x[:, :64]
|
||||
for d in policy_cfg.hiddens:
|
||||
x = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x))
|
||||
x = nn.Dense(self.num_actions, kernel_init=kernel_init)(x)
|
||||
log_std = self.param("log_std",
|
||||
lambda _: jnp.full(self.num_actions, float(policy_cfg.initial_log_std)))
|
||||
return x, log_std, {}
|
||||
|
||||
class Value(DeterministicMixin, Model):
|
||||
def __init__(self, observation_space, action_space, device=None, **kwargs):
|
||||
Model.__init__(self, observation_space, action_space, device, **kwargs)
|
||||
DeterministicMixin.__init__(self, value_cfg.clip_actions)
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, inputs, role):
|
||||
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
|
||||
x = inputs["states"]
|
||||
# Critic: obs(45) + base_vel(3) + heights(187) = 235
|
||||
# Layout: [code(19) | obs(45) | base_vel(3) | heights(187)]
|
||||
x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1)
|
||||
for d in value_cfg.hiddens:
|
||||
x_c = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x_c))
|
||||
return nn.Dense(1, kernel_init=kernel_init)(x_c), {}
|
||||
|
||||
models = {
|
||||
"policy": Policy(wrapper.observation_space, wrapper.action_space, None),
|
||||
"value": Value(wrapper.observation_space, wrapper.action_space, None),
|
||||
}
|
||||
for role, m in models.items():
|
||||
m.init_state_dict(role)
|
||||
|
||||
# Agent
|
||||
ppo_cfg = cfg.runner.agent.to_dict()
|
||||
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
|
||||
_add_runtime_config(ppo_cfg, wrapper, log_dir=log_dir)
|
||||
|
||||
mem_size = cfg.runner.memory.memory_size
|
||||
if mem_size == -1:
|
||||
mem_size = ppo_cfg.get("rollouts", 24)
|
||||
memory = RandomMemory(memory_size=mem_size, num_envs=cfg.num_envs, device=None)
|
||||
|
||||
agent = DreamWaQAgent(
|
||||
models=models, memory=memory, cfg=ppo_cfg,
|
||||
observation_space=wrapper.observation_space,
|
||||
action_space=wrapper.action_space, device=None,
|
||||
cenet=cenet, cenet_params=cenet_params, wrapper=wrapper,
|
||||
kl_weight=1.0, vae_lr=1e-3, save_dir=log_dir,
|
||||
)
|
||||
|
||||
trainer_cfg = {"timesteps": cfg.runner.trainer.timesteps, "headless": True}
|
||||
SequentialTrainer(cfg=trainer_cfg, env=wrapper, agents=agent).train()
|
||||
|
||||
# Save final VAE params
|
||||
pickle.dump(agent._cenet_params, open(f"{log_dir}/cenet_params.pkl", "wb"))
|
||||
print(f"[DreamWaQ] CENet weights saved → {log_dir}")
|
||||
|
||||
def play(self, policy_path=None):
|
||||
"""Render env with trained policy (uses wrapper + CENet)."""
|
||||
from motrix_rl.skrl.jax.wrap_np import SkrlNpWrapper as RenderWrapper
|
||||
cfg = self._rlcfg
|
||||
raw_env = env_registry.make(self._env_name, num_envs=min(64, cfg.play_num_envs))
|
||||
env = RenderWrapper(raw_env) # adds render() method
|
||||
env_cfg = raw_env._cfg
|
||||
|
||||
import msgpack
|
||||
rng = jax.random.PRNGKey(cfg.runner.seed)
|
||||
rng, crng = jax.random.split(rng)
|
||||
cenet = CENet()
|
||||
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
|
||||
|
||||
# Load CENet params
|
||||
if policy_path and os.path.exists(policy_path):
|
||||
# Load full agent checkpoint and extract
|
||||
with open(policy_path, 'rb') as f:
|
||||
ckpt = pickle.load(f)
|
||||
# Load policy params
|
||||
policy_raw = msgpack.unpackb(ckpt['policy'])
|
||||
policy_params = {}
|
||||
for name, val in policy_raw['params'].items():
|
||||
if isinstance(val, dict):
|
||||
policy_params[name] = {
|
||||
k: _decode_msgpack_array(v) for k, v in val.items()
|
||||
}
|
||||
else:
|
||||
policy_params[name] = _decode_msgpack_array(val)
|
||||
|
||||
# Try loading VAE from base skrl dir
|
||||
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
|
||||
vae_path = os.path.join(log_dir, "cenet_params.pkl")
|
||||
if not os.path.exists(vae_path):
|
||||
# Find latest vae_N.pkl
|
||||
vae_files = sorted([f for f in os.listdir(log_dir) if f.startswith("vae_")])
|
||||
if vae_files:
|
||||
vae_path = os.path.join(log_dir, vae_files[-1])
|
||||
if os.path.exists(vae_path):
|
||||
with open(vae_path, 'rb') as f:
|
||||
cenet_params = pickle.load(f)
|
||||
print(f"[Play] VAE loaded: {vae_path}")
|
||||
else:
|
||||
print("[Play] WARNING: VAE params not found, using random init")
|
||||
cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))
|
||||
else:
|
||||
print("[Play] No policy, random actions")
|
||||
policy_params = None
|
||||
cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))
|
||||
|
||||
wrapper = DreamWaQWrapper(raw_env, cenet, cenet_params,
|
||||
rng=jax.random.PRNGKey(cfg.runner.seed + 1))
|
||||
raw_env.init_state()
|
||||
|
||||
import time
|
||||
n = raw_env._num_envs
|
||||
print(f"[Play] {n} envs, press Ctrl+C to stop")
|
||||
try:
|
||||
while True:
|
||||
if policy_params is not None:
|
||||
aug_obs = wrapper._augment(raw_env._state.obs)
|
||||
actions = np.array(_policy_forward(jnp.array(aug_obs), policy_params))
|
||||
else:
|
||||
actions = np.zeros((n, 12), dtype=np.float32)
|
||||
|
||||
wrapper.step(actions)
|
||||
env.render()
|
||||
time.sleep(0.01)
|
||||
except KeyboardInterrupt:
|
||||
env.close()
|
||||
print("[Play] Stopped")
|
||||
|
||||
|
||||
def _decode_msgpack_array(ext) -> np.ndarray | None:
|
||||
"""Decode flax-serialized msgpack ExtType."""
|
||||
import msgpack
|
||||
if not hasattr(ext, "code"): return None
|
||||
parts = msgpack.unpackb(ext.data, raw=False)
|
||||
if not isinstance(parts, list) or len(parts) < 3: return None
|
||||
shape = []
|
||||
def _flatten(s):
|
||||
if isinstance(s, list):
|
||||
for x in s: _flatten(x)
|
||||
elif isinstance(s, int): shape.append(s)
|
||||
_flatten(parts[0])
|
||||
return np.frombuffer(parts[2], dtype=np.dtype(parts[1])).reshape(shape)
|
||||
|
||||
|
||||
def _policy_forward(x, params):
|
||||
"""Forward pass through policy (first 64 dims only, apply elu)."""
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
w0, b0 = jnp.array(params['Dense_0']['kernel']), jnp.array(params['Dense_0']['bias'])
|
||||
w1, b1 = jnp.array(params['Dense_1']['kernel']), jnp.array(params['Dense_1']['bias'])
|
||||
w2, b2 = jnp.array(params['Dense_2']['kernel']), jnp.array(params['Dense_2']['bias'])
|
||||
w3, b3 = jnp.array(params['Dense_3']['kernel']), jnp.array(params['Dense_3']['bias'])
|
||||
x = x[:, :64] # code(19) + obs(45)
|
||||
x = jax.nn.elu(x @ w0 + b0)
|
||||
x = jax.nn.elu(x @ w1 + b1)
|
||||
x = jax.nn.elu(x @ w2 + b2)
|
||||
return x @ w3 + b3
|
||||
@@ -185,6 +185,7 @@ class Trainer:
|
||||
sim_backend: str = None,
|
||||
enable_render: bool = False,
|
||||
cfg_override: dict = None,
|
||||
env_cfg_override: dict = None,
|
||||
) -> None:
|
||||
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
|
||||
if cfg_override is not None:
|
||||
@@ -193,13 +194,18 @@ class Trainer:
|
||||
self._env_name = env_name
|
||||
self._sim_backend = sim_backend
|
||||
self._enable_render = enable_render
|
||||
self._env_cfg_override = env_cfg_override
|
||||
|
||||
def train(self) -> None:
|
||||
def train(self, checkpoint: str = None) -> None:
|
||||
"""
|
||||
Start training the agent.
|
||||
|
||||
Args:
|
||||
checkpoint: Optional path to a checkpoint to resume from.
|
||||
"""
|
||||
rlcfg = self._rlcfg
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
|
||||
num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
|
||||
|
||||
set_seed(rlcfg.runner.seed)
|
||||
skrl_env = wrap_env(env, self._enable_render)
|
||||
@@ -209,6 +215,12 @@ class Trainer:
|
||||
# Add runtime-specific configuration
|
||||
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
|
||||
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
|
||||
|
||||
if checkpoint:
|
||||
agent.load(checkpoint)
|
||||
logger = __import__("logging").getLogger(__name__)
|
||||
logger.info(f"Resumed from checkpoint: {checkpoint}")
|
||||
|
||||
cfg_trainer = {
|
||||
"timesteps": rlcfg.runner.trainer.timesteps,
|
||||
"headless": not self._enable_render,
|
||||
@@ -220,7 +232,17 @@ class Trainer:
|
||||
import time
|
||||
|
||||
rlcfg = self._rlcfg
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
|
||||
# Enable play_mode + merge constructor-level overrides (e.g. force_phase)
|
||||
play_override = {"play_mode": True}
|
||||
if self._env_cfg_override:
|
||||
play_override.update(self._env_cfg_override)
|
||||
try:
|
||||
env = env_registry.make(
|
||||
self._env_name, sim_backend=self._sim_backend,
|
||||
num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
|
||||
)
|
||||
except ValueError:
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
|
||||
|
||||
set_seed(rlcfg.runner.seed)
|
||||
env = wrap_env(env, self._enable_render)
|
||||
|
||||
@@ -185,6 +185,7 @@ class Trainer:
|
||||
sim_backend: str = None,
|
||||
enable_render: bool = False,
|
||||
cfg_override: dict = None,
|
||||
env_cfg_override: dict = None,
|
||||
) -> None:
|
||||
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="torch")
|
||||
if cfg_override is not None:
|
||||
@@ -193,13 +194,18 @@ class Trainer:
|
||||
self._env_name = env_name
|
||||
self._sim_backend = sim_backend
|
||||
self._enable_render = enable_render
|
||||
self._env_cfg_override = env_cfg_override
|
||||
|
||||
def train(self) -> None:
|
||||
def train(self, checkpoint: str = None) -> None:
|
||||
"""
|
||||
Start training the agent.
|
||||
|
||||
Args:
|
||||
checkpoint: Optional path to a checkpoint to resume from.
|
||||
"""
|
||||
rlcfg = self._rlcfg
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
|
||||
num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
|
||||
set_seed(rlcfg.runner.seed)
|
||||
skrl_env = wrap_env(env, self._enable_render)
|
||||
models = self._make_model(skrl_env, rlcfg)
|
||||
@@ -208,6 +214,12 @@ class Trainer:
|
||||
# Add runtime-specific configuration
|
||||
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
|
||||
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
|
||||
|
||||
if checkpoint:
|
||||
agent.load(checkpoint)
|
||||
logger = __import__("logging").getLogger(__name__)
|
||||
logger.info(f"Resumed from checkpoint: {checkpoint}")
|
||||
|
||||
cfg_trainer = {
|
||||
"timesteps": rlcfg.runner.trainer.timesteps,
|
||||
"headless": not self._enable_render,
|
||||
@@ -219,7 +231,17 @@ class Trainer:
|
||||
import time
|
||||
|
||||
rlcfg = self._rlcfg
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
|
||||
# Enable play_mode if the env config supports it (for multi-terrain random spawn)
|
||||
play_override = {"play_mode": True}
|
||||
if self._env_cfg_override:
|
||||
play_override.update(self._env_cfg_override)
|
||||
try:
|
||||
env = env_registry.make(
|
||||
self._env_name, sim_backend=self._sim_backend,
|
||||
num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
|
||||
)
|
||||
except ValueError:
|
||||
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
|
||||
set_seed(rlcfg.runner.seed)
|
||||
env = wrap_env(env, self._enable_render)
|
||||
models = self._make_model(env, rlcfg)
|
||||
|
||||
@@ -135,7 +135,7 @@ class rslrl:
|
||||
runner.actor.class_name = (
|
||||
"motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel")
|
||||
runner.actor.hidden_dims = [512, 256, 128]
|
||||
runner.actor.init_noise_std = 1.0
|
||||
runner.actor.init_noise_std = 0.5 # 降噪声先学站稳，再探索行走
|
||||
|
||||
# Critic：标准 MLPModel，输入 privileged_obs
|
||||
runner.critic.class_name = "MLPModel"
|
||||
|
||||
42
motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py
Normal file
42
motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""DreamWaQ training config — Manaro-Alpha aligned.
|
||||
|
||||
Actor [512,256,128], Critic [512,256,128], CENet [128,64].
|
||||
lr=1e-3, vae_lr=1e-3, KL_weight=1.0, entropy=0.01.
|
||||
"""
|
||||
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401
|
||||
from motrix_rl.registry import rlcfg
|
||||
from motrix_rl.skrl.config import SkrlCfg
|
||||
|
||||
|
||||
def _apply(cfg: SkrlCfg):
|
||||
# PPO hyperparams (exact DreamWaQ match)
|
||||
cfg.runner.agent.rollouts = 24 # num_steps_per_env
|
||||
cfg.runner.agent.learning_epochs = 5 # num_learning_epochs
|
||||
cfg.runner.agent.mini_batches = 4 # num_mini_batches
|
||||
cfg.runner.agent.learning_rate = 1e-3
|
||||
cfg.runner.agent.entropy_loss_scale = 0.01
|
||||
cfg.runner.agent.ratio_clip = 0.2
|
||||
cfg.runner.agent.discount_factor = 0.99
|
||||
cfg.runner.agent.lam = 0.95
|
||||
|
||||
# Network (DreamWaQ: [512,256,128])
|
||||
cfg.runner.models.policy.hiddens = [512, 256, 128]
|
||||
cfg.runner.models.policy.hidden_activation = ["elu"]
|
||||
# initial_log_std: upstream uses 0.0 (std=1) in IsaacGym, but in MotrixSim with
|
||||
# action_scale=0.25 that gives ±0.25 rad jitter/step → robot falls immediately and
|
||||
# never learns to stand. -1.0 (std=0.37, ±0.09 rad) lets it stay upright early and
|
||||
# bootstrap walking. (MotrixSim<->IsaacGym reality-gap adjustment.)
|
||||
cfg.runner.models.policy.initial_log_std = -1.0
|
||||
cfg.runner.models.value.hiddens = [512, 256, 128]
|
||||
cfg.runner.models.value.hidden_activation = ["elu"]
|
||||
|
||||
cfg.runner.memory.memory_size = 24
|
||||
cfg.runner.agent.experiment.checkpoint_interval = 200
|
||||
cfg.runner.agent.experiment.write_interval = 200
|
||||
|
||||
|
||||
@rlcfg("go1-dreamwaq-walk")
|
||||
class DreamWaQ(SkrlCfg):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
_apply(self)
|
||||
64
motrix_rl/src/motrix_rl/tasks/go1_go2style.py
Normal file
64
motrix_rl/src/motrix_rl/tasks/go1_go2style.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""Training configs for Go1 go2style environments.
|
||||
|
||||
Importing triggers registration of: flat, stairs, combined.
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
|
||||
import motrix_envs.locomotion.go1.go2style_flat # noqa: F401
|
||||
import motrix_envs.locomotion.go1.go2style_stairs # noqa: F401
|
||||
import motrix_envs.locomotion.go1.go2style_combined # noqa: F401
|
||||
import motrix_envs.locomotion.go1.go2style_cts # noqa: F401
|
||||
import motrix_envs.locomotion.go1.orig_combined # noqa: F401
|
||||
|
||||
from motrix_rl.registry import rlcfg
|
||||
from motrix_rl.skrl.config import SkrlCfg
|
||||
|
||||
|
||||
def _apply_go2style_ppo(cfg: SkrlCfg):
|
||||
"""Apply go2_rl_gym-style PPO params to config."""
|
||||
cfg.runner.models.policy.hiddens = [512, 256, 128]
|
||||
cfg.runner.models.value.hiddens = [512, 256, 128]
|
||||
cfg.runner.agent.rollouts = 24
|
||||
cfg.runner.agent.learning_epochs = 5
|
||||
cfg.runner.agent.mini_batches = 4
|
||||
cfg.runner.agent.learning_rate = 1e-3
|
||||
cfg.runner.trainer.timesteps = 60000
|
||||
|
||||
|
||||
def _apply_orig_ppo(cfg: SkrlCfg):
|
||||
"""Match original Go1 stairs config: [512,256,128], 30k steps."""
|
||||
cfg.runner.models.policy.hiddens = [512, 256, 128]
|
||||
cfg.runner.models.value.hiddens = [512, 256, 128]
|
||||
cfg.runner.agent.rollouts = 24
|
||||
cfg.runner.agent.learning_epochs = 5
|
||||
cfg.runner.agent.mini_batches = 4
|
||||
cfg.runner.agent.learning_rate = 1e-3
|
||||
cfg.runner.trainer.timesteps = 30000 # same as original stairs
|
||||
|
||||
|
||||
class skrl:
|
||||
@rlcfg("go1-flat-terrain-walk-go2style")
|
||||
@dataclass
|
||||
class Flat(SkrlCfg):
|
||||
def __post_init__(self): _apply_go2style_ppo(self)
|
||||
|
||||
@rlcfg("go1-stairs-terrain-walk-go2style")
|
||||
@dataclass
|
||||
class Stairs(SkrlCfg):
|
||||
def __post_init__(self): _apply_go2style_ppo(self)
|
||||
|
||||
@rlcfg("go1-combined-terrain-walk-go2style")
|
||||
@dataclass
|
||||
class Combined(SkrlCfg):
|
||||
def __post_init__(self): _apply_go2style_ppo(self)
|
||||
|
||||
@rlcfg("go1-cts-flat-walk-go2style")
|
||||
@dataclass
|
||||
class CTS(SkrlCfg):
|
||||
def __post_init__(self): _apply_go2style_ppo(self)
|
||||
|
||||
@rlcfg("go1-orig-combined-walk")
|
||||
@dataclass
|
||||
class OrigCombined(SkrlCfg):
|
||||
"""Original Go1 (PD 80, 45-dim, [256,128,64]) on flat+stairs."""
|
||||
def __post_init__(self): _apply_orig_ppo(self)
|
||||
Reference in New Issue
Block a user