| """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) |
|
|
