199 lines
9.4 KiB
Python
199 lines
9.4 KiB
Python
"""Differentiable rewards, full-return actor objective and TD(lambda) critic targets."""
|
|
|
|
import math
|
|
from contextlib import contextmanager
|
|
import torch
|
|
from .world_model import GRAVITY, LIN_VEL, ANG_VEL, LEG_POS, LEG_VEL, WHEEL_VEL
|
|
from .policy import frame, push_history, critic_input
|
|
|
|
|
|
@contextmanager
|
|
def frozen(*modules):
|
|
"""Freeze weights without disabling gradients through inputs/actions."""
|
|
parameters = [(p, p.requires_grad) for m in modules for p in m.parameters()]
|
|
try:
|
|
for p, _ in parameters:
|
|
p.requires_grad_(False)
|
|
yield
|
|
finally:
|
|
for p, enabled in parameters:
|
|
p.requires_grad_(enabled)
|
|
|
|
|
|
def survival(state):
|
|
gravity = torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
|
|
# Smooth proxy for the environment's 70-degree fall cutoff, width ~5 deg.
|
|
return torch.sigmoid((-gravity[:, 2] - math.cos(math.radians(70))) / 0.08)
|
|
|
|
|
|
def continuation_value(state, constants):
|
|
"""Imagined continuation probability for the current task.
|
|
|
|
The v4 proxy treats a 70-degree tilt as an immediate fall. The upward task
|
|
instead terminates only after 4 s of bad orientation, which is 200 control
|
|
steps, far beyond the 32/64-step imagination horizon, and the counter is not
|
|
part of the state. Applying the v4 proxy there truncates the return after
|
|
the first step whenever the robot is inverted, so the actor cannot learn to
|
|
flip back.
|
|
"""
|
|
if constants.get("upright_mode") == "upward":
|
|
return torch.ones(state.shape[0], dtype=state.dtype, device=state.device)
|
|
return survival(state)
|
|
|
|
|
|
def reward_terms(state, action, previous, command, constants, survival_prob, scan=None):
|
|
v, w = state[:, LIN_VEL], state[:, ANG_VEL]
|
|
q, qd, wd = state[:, LEG_POS], state[:, LEG_VEL], state[:, WHEEL_VEL]
|
|
standing = (command.square().sum(-1) < 1e-6).float()
|
|
# Endpoint PD torque surrogate: actual MuJoCo forces also reflect implicit
|
|
# integration over substeps. The low-weight effort term is approximate.
|
|
leg_torque = (constants["kp"] * (constants.get("leg_action_scale", 0.25) * action[:, :12] - q)
|
|
- constants["kd"] * qd).clamp(-constants["effort"], constants["effort"])
|
|
wheel_limit = constants.get("wheel_velocity_limit", 35.)
|
|
wheel_effort = constants.get("wheel_effort", 45.)
|
|
wheel_torque = (constants.get("wheel_kd", 2.) * ((constants.get("wheel_action_scale", 15.) * action[:, 12:]).clamp(-wheel_limit, wheel_limit) - wd)).clamp(-wheel_effort, wheel_effort)
|
|
absolute_q = q + constants["default_q"]
|
|
limits = ((constants["limits"][..., 0] - absolute_q).clamp_min(0)
|
|
+ (absolute_q - constants["limits"][..., 1]).clamp_min(0)).sum(-1)
|
|
posture = q.square().sum(-1)
|
|
if constants.get("terrain_posture_v2", False) and not constants.get("wheel_local_rewards", False):
|
|
if scan is None:
|
|
raise ValueError("Terrain v2 posture reward requires the privileged scan")
|
|
from .terrain_posture import posture_cost
|
|
posture = posture_cost(q, scan, constants.get("posture_relief_floor", 0.125))
|
|
terrain = constants.get("terrain_rewards", False)
|
|
gravity = state[:, GRAVITY]
|
|
if constants.get("upright_mode") == "upward":
|
|
# Sign-aware recovery reward: 1 upright, 0.5 sideways, 0 upside down.
|
|
upright = (0.5 * (1. - gravity[:, 2])).clamp(0., 1.)
|
|
else:
|
|
upright = torch.exp(
|
|
-gravity[:, :2].square().sum(-1) / constants.get("upright_std", 0.45)**2)
|
|
terms = {
|
|
"track_linear_velocity": torch.exp(-((v[:, :2] - command[:, :2]).square().sum(-1)
|
|
+ (0. if terrain else v[:, 2].square())) / 0.25),
|
|
"track_angular_velocity": torch.exp(-((w[:, 2] - command[:, 2]).square()
|
|
+ (0. if terrain else w[:, :2].square().sum(-1))) / 0.5),
|
|
"upright": upright,
|
|
"leg_posture": posture,
|
|
"standstill": standing * (v.square().sum(-1) + w.square().sum(-1)
|
|
+ 0.01 * wd.square().sum(-1)),
|
|
"action_rate_l2": (action - previous).square().sum(-1),
|
|
"joint_torques_l2": leg_torque.square().sum(-1) + wheel_torque.square().sum(-1),
|
|
"dof_pos_limits": limits,
|
|
"terminated": 1. - survival_prob,
|
|
}
|
|
|
|
if constants.get("rolling_rewards", False) and not constants.get("wheel_local_rewards", False):
|
|
if scan is None:
|
|
raise ValueError("Rolling rewards require a terrain scan")
|
|
from .rolling_rewards import rolling_costs
|
|
terms.update(rolling_costs(state, command, scan, constants["default_q"]))
|
|
if constants.get("obstacle_progress", False):
|
|
from .obstacle_rewards import obstacle_progress
|
|
terms["obstacle_progress"] = obstacle_progress(state, command, scan)
|
|
if constants.get("wheel_body_intrusion", False):
|
|
from .wheel_intrusion import wheel_body_intrusion
|
|
terms["wheel_body_intrusion"] = wheel_body_intrusion(
|
|
state, constants["default_q"])
|
|
if constants.get("wheel_local_rewards", False):
|
|
from .wheel_terrain_rewards import wheel_terms
|
|
terms.update(wheel_terms(state, command, scan, constants["default_q"]))
|
|
if constants.get("bad_orientation_penalty", False):
|
|
from .upward import bad_orientation_value
|
|
terms["bad_orientation_penalty"] = bad_orientation_value(gravity)
|
|
from .upward import apply_reward_clamps
|
|
apply_reward_clamps(terms, constants)
|
|
return terms
|
|
|
|
|
|
def lambda_returns(rewards, continuation, next_values, gamma=0.99, lam=0.95):
|
|
carry = next_values[-1]
|
|
result = []
|
|
for t in reversed(range(len(rewards))):
|
|
carry = rewards[t] + gamma * continuation[t] * ((1 - lam) * next_values[t] + lam * carry)
|
|
result.append(carry)
|
|
return torch.stack(result[::-1])
|
|
|
|
|
|
def discounted_objective(rewards, continuation, terminal_value, gamma=.99):
|
|
alive = torch.ones_like(terminal_value)
|
|
total = torch.zeros_like(terminal_value)
|
|
discount = 1.
|
|
for reward, cont in zip(rewards, continuation):
|
|
total = total + discount * alive * reward
|
|
alive = alive * cont
|
|
discount *= gamma
|
|
return total + discount * alive * terminal_value
|
|
|
|
|
|
def make_imagined_step(model, actor, constants, weights, env_action_order,
|
|
dt=.02, noise=True, field=None):
|
|
"""Pure single-step graph, shared by eager and compiled execution."""
|
|
def step(state, previous, history, pose, scan, command, context, bias, timers):
|
|
action = actor(history)
|
|
model_context = context if field is None else torch.cat((context, scan), -1)
|
|
collision_prob = None
|
|
if constants.get('nonwheel_collision', False):
|
|
nxt, delta, _, collision_logit = model.forward_all(state, action, previous, model_context)
|
|
collision_prob = collision_logit.squeeze(-1).sigmoid()
|
|
else:
|
|
nxt, delta = model(state, action, previous, model_context)
|
|
if field is not None:
|
|
from .terrain import advance_pose
|
|
pose = advance_pose(pose, delta)
|
|
next_scan = None if field is None else field.scan(nxt, pose)
|
|
cont = continuation_value(nxt, constants)
|
|
if collision_prob is not None:
|
|
cont = cont * (1. - collision_prob)
|
|
terms = reward_terms(nxt, action, previous, command, constants, cont,
|
|
scan=next_scan)
|
|
if constants.get('pair_body_intrusion', False):
|
|
from .wheel_intrusion import pair_body_intrusion
|
|
terms['pair_body_intrusion'], timers = pair_body_intrusion(
|
|
nxt, constants['default_q'], timers, dt)
|
|
reward = dt * sum(weights[name] * value for name, value in terms.items())
|
|
latest = frame(nxt, action[:, env_action_order], command, bias, noise=noise,
|
|
wheel_vel=actor.actor_wheel_vel)
|
|
history = push_history(history, latest, actor.history, actor.terms)
|
|
return nxt, action, history, pose, next_scan, reward, cont, timers
|
|
return step
|
|
|
|
|
|
def rollout(model, actor, target_critic, seed, constants, weights, env_action_order,
|
|
horizon=16, dt=0.02, gamma=0.99, lam=0.95, noise=True, field=None,
|
|
step_fn=None):
|
|
"""Actor: full discounted return + terminal value. Critic: detached TD(lambda).
|
|
|
|
Caller freezes dynamics and target critic weights, preserving state gradients.
|
|
Reuse one terrain scan at each state, including the rollout's initial state.
|
|
"""
|
|
state, previous, history = seed["state"], seed["previous"], seed["history"]
|
|
command, context, bias = seed["command"], seed["context"], seed["bias"]
|
|
pose = seed.get("world_pose")
|
|
timers = seed.get("pair_timer")
|
|
if timers is None:
|
|
timers = torch.zeros(state.shape[0], 4, device=state.device)
|
|
def value_input(s, prev, scan):
|
|
base = critic_input(s, command, context, prev)
|
|
return base if field is None else torch.cat((base, scan), -1)
|
|
if step_fn is None:
|
|
step_fn = make_imagined_step(model, actor, constants, weights, env_action_order,
|
|
dt, noise, field)
|
|
rewards, continuations, inputs, next_inputs = [], [], [], []
|
|
scan = None if field is None else field.scan(state, pose)
|
|
for _ in range(horizon):
|
|
inputs.append(value_input(state, previous, scan).detach())
|
|
state, previous, history, pose, scan, reward, cont, timers = step_fn(
|
|
state, previous, history, pose, scan, command, context, bias, timers)
|
|
rewards.append(reward)
|
|
continuations.append(cont)
|
|
next_inputs.append(value_input(state, previous, scan).detach())
|
|
final_value = target_critic(value_input(state, previous, scan)).squeeze(-1)
|
|
r, c = torch.stack(rewards), torch.stack(continuations)
|
|
objective = discounted_objective(r, c, final_value, gamma)
|
|
with torch.no_grad():
|
|
next_values = target_critic(torch.stack(next_inputs)).squeeze(-1)
|
|
targets = lambda_returns(r.detach(), c.detach(), next_values, gamma, lam)
|
|
return -objective.mean(), torch.stack(inputs), targets
|