File size: 1,517 Bytes
d8290eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | """Identical collector/learner action conversion in actual joint radians."""
from __future__ import annotations
import numpy as np
from .data import ActionContract, BASE_ACTION_DIM, RESIDUAL_ACTION_DIM, vector
class ResidualActionAdapter:
def __init__(self, contract: ActionContract):
self.contract = contract
self.indices = contract.indices
self.scales_rad = np.array(contract.residual_scales_rad, dtype=np.float32)
self.base_scales = np.array(contract.action_scale, dtype=np.float32)
self.action_increments = self.scales_rad / self.base_scales[self.indices]
def combine(self, base_action, residual):
base = np.clip(vector(base_action, BASE_ACTION_DIM), -10.0, 10.0)
residual = np.clip(vector(residual, RESIDUAL_ACTION_DIM), -1.0, 1.0)
combined = base.copy()
combined[self.indices] += self.action_increments * residual
combined = np.clip(combined, -10.0, 10.0)
executed_rad = (combined[self.indices] - base[self.indices]) * self.base_scales[self.indices]
return combined, executed_rad
def combine_tensor(self, base_action, residual):
import torch
indices = torch.as_tensor(self.indices, device=base_action.device)
scales = torch.as_tensor(self.action_increments, device=base_action.device)
combined = base_action.clamp(-10.0, 10.0).clone()
combined[:, indices] = combined[:, indices] + residual.clamp(-1.0, 1.0) * scales
return combined.clamp(-10.0, 10.0)
|