Add files using upload-large-folder tool

7985e7e verified 18 days ago

7.16 kB

"""
ProxyValueModel: batched differentiable value function for gradient guidance.
value = objective_weight * objective
      - obstacle_weight * obstacle_penalty
      - dynamics_weight * dynamics_penalty
All operations are in PyTorch for autograd compatibility.
"""

import numpy as np
import torch
from torch import nn

from obstacles import PILLARS, NOVEL_CIRCULAR, NOVEL_PLANAR


class ProxyValueModel(nn.Module):
    """Differentiable value function for trajectory guidance.
    Evaluates on normalized trajectories of shape (B, H, action_dim + state_dim).
    Internally unnormalizes positions using the dataset normalizer bounds.
    """

    def __init__(
        self,
        horizon,
        action_dim,
        state_dim,
        normalizer,
        objective="",
        constraints=False,
        dynamics=False,
        objective_weight=1.0,
        obstacle_weight=100.0,
        dynamics_weight=0.0,
        obstacle_margin=0.02,
        dynamics_model=None,
        y_goal=0.35,
    ):
        super().__init__()
        self.horizon = horizon
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.objective = objective
        self.constraints = constraints
        self.dynamics = dynamics
        self.objective_weight = objective_weight
        self.obstacle_weight = obstacle_weight
        self.dynamics_weight = dynamics_weight
        self.obstacle_margin = obstacle_margin
        self.y_goal = y_goal

        # normalizer bounds for observations (x_pos=dim2, y_pos=dim3)
        obs_norm = normalizer.normalizers["observations"]
        self.register_buffer("obs_mins", torch.tensor(obs_norm.mins, dtype=torch.float32))
        self.register_buffer("obs_maxs", torch.tensor(obs_norm.maxs, dtype=torch.float32))

        # dynamics model (A, B, c) as buffers
        if dynamics and dynamics_model is not None:
            self.register_buffer("A", torch.tensor(dynamics_model["A"], dtype=torch.float32))
            self.register_buffer("B", torch.tensor(dynamics_model["B"], dtype=torch.float32))
            self.register_buffer("c", torch.tensor(dynamics_model["c"], dtype=torch.float32).squeeze())

        # precompute obstacle tensors
        if constraints:
            circular = PILLARS + NOVEL_CIRCULAR
            centers = torch.tensor([obs.center for obs in circular], dtype=torch.float32)
            radii = torch.tensor([obs.radius for obs in circular], dtype=torch.float32)
            self.register_buffer("obs_centers", centers)
            self.register_buffer("obs_radii", radii)

            if NOVEL_PLANAR:
                slopes = torch.tensor([obs.slope for obs in NOVEL_PLANAR], dtype=torch.float32)
                intercepts = torch.tensor([obs.intercept for obs in NOVEL_PLANAR], dtype=torch.float32)
                self.register_buffer("planar_slopes", slopes)
                self.register_buffer("planar_intercepts", intercepts)
            else:
                self.planar_slopes = None

    def _unnorm_obs(self, obs_norm):
        """Unnormalize observations from [-1,1] to real-world coords.
        Args:
            obs_norm: (..., state_dim) normalized observations
        Returns:
            (..., state_dim) real-world observations
        """
        return (obs_norm + 1.0) / 2.0 * (self.obs_maxs - self.obs_mins) + self.obs_mins

    def forward(self, x):
        """Compute value for a batch of trajectories.
        Args:
            x: (B, H, action_dim + state_dim) normalized trajectories
        Returns:
            value: (B,) scalar value per trajectory (higher is better)
        """
        B = x.shape[0]
        actions = x[:, :, :self.action_dim]          # (B, H, action_dim)
        obs_norm = x[:, :, self.action_dim:]          # (B, H, state_dim)
        obs_real = self._unnorm_obs(obs_norm)         # (B, H, state_dim)

        value = torch.zeros(B, device=x.device)

        # --- objective ---
        if self.objective == "distance":
            # terminal y-position cost: -(y_terminal - y_goal)^2
            y_terminal = obs_real[:, -1, 3]  # (B,) — y_pos is dim 3
            value = value + self.objective_weight * (-(y_terminal - self.y_goal) ** 2)

        # --- obstacle penalty ---
        if self.constraints and self.obstacle_weight > 0:
            value = value - self.obstacle_weight * self._obstacle_penalty(obs_real)

        # --- dynamics penalty ---
        if self.dynamics and self.dynamics_weight > 0:
            value = value - self.dynamics_weight * self._dynamics_penalty(actions, obs_norm)

        return value

    def _obstacle_penalty(self, obs_real):
        """Compute squared obstacle violation penalty.
        Args:
            obs_real: (B, H, state_dim) real-world observations
        Returns:
            penalty: (B,) sum of squared violations
        """
        x_pos = obs_real[:, :, 2]  # (B, H)
        y_pos = obs_real[:, :, 3]  # (B, H)

        penalty = torch.zeros(obs_real.shape[0], device=obs_real.device)

        # circular obstacles: violation = clamp(r + margin - dist, min=0)
        # obs_centers: (N_obs, 2), obs_radii: (N_obs,)
        for i in range(self.obs_centers.shape[0]):
            cx, cy = self.obs_centers[i]
            r = self.obs_radii[i]
            dist = torch.sqrt((x_pos - cx) ** 2 + (y_pos - cy) ** 2)  # (B, H)
            violation = torch.clamp(r + self.obstacle_margin - dist, min=0)  # (B, H)
            penalty = penalty + (violation ** 2).sum(dim=1)  # (B,)

        # planar obstacles: violation = clamp(-(slope * x + intercept - y) + margin, min=0)
        # feasible: slope * x + intercept - y >= 0
        if self.planar_slopes is not None:
            for i in range(self.planar_slopes.shape[0]):
                s = self.planar_slopes[i]
                b = self.planar_intercepts[i]
                feasibility = s * x_pos + b - y_pos  # (B, H) — positive = feasible
                violation = torch.clamp(self.obstacle_margin - feasibility, min=0)  # (B, H)
                penalty = penalty + (violation ** 2).sum(dim=1)

        return penalty

    def _dynamics_penalty(self, actions, obs_norm):
        """Compute squared dynamics residual penalty in normalized space.
        s_{t+1} = A @ s_t + B @ a_t + c  (fitted linear model)
        Args:
            actions: (B, H, action_dim) normalized actions
            obs_norm: (B, H, state_dim) normalized observations
        Returns:
            penalty: (B,) sum of squared residuals
        """
        # s_t for t=0..H-2, s_{t+1} for t=1..H-1
        s_t = obs_norm[:, :-1, :]      # (B, H-1, state_dim)
        s_next = obs_norm[:, 1:, :]    # (B, H-1, state_dim)
        a_t = actions[:, :-1, :]       # (B, H-1, action_dim)

        # predicted: A @ s_t + B @ a_t + c
        # A: (state_dim, state_dim), B: (state_dim, action_dim), c: (state_dim,)
        pred = (torch.einsum("ij,bhj->bhi", self.A, s_t)
                + torch.einsum("ij,bhj->bhi", self.B, a_t)
                + self.c)  # (B, H-1, state_dim)

        residual = pred - s_next  # (B, H-1, state_dim)
        penalty = (residual ** 2).sum(dim=(1, 2))  # (B,)
        return penalty

	"""
	ProxyValueModel: batched differentiable value function for gradient guidance.

	value = objective_weight * objective
	- obstacle_weight * obstacle_penalty
	- dynamics_weight * dynamics_penalty

	All operations are in PyTorch for autograd compatibility.
	"""

	import numpy as np
	import torch
	from torch import nn

	from obstacles import PILLARS, NOVEL_CIRCULAR, NOVEL_PLANAR


	class ProxyValueModel(nn.Module):
	"""Differentiable value function for trajectory guidance.

	Evaluates on normalized trajectories of shape (B, H, action_dim + state_dim).
	Internally unnormalizes positions using the dataset normalizer bounds.
	"""

	def __init__(
	self,
	horizon,
	action_dim,
	state_dim,
	normalizer,
	objective="",
	constraints=False,
	dynamics=False,
	objective_weight=1.0,
	obstacle_weight=100.0,
	dynamics_weight=0.0,
	obstacle_margin=0.02,
	dynamics_model=None,
	y_goal=0.35,
	):
	super().__init__()
	self.horizon = horizon
	self.action_dim = action_dim
	self.state_dim = state_dim
	self.objective = objective
	self.constraints = constraints
	self.dynamics = dynamics
	self.objective_weight = objective_weight
	self.obstacle_weight = obstacle_weight
	self.dynamics_weight = dynamics_weight
	self.obstacle_margin = obstacle_margin
	self.y_goal = y_goal

	# normalizer bounds for observations (x_pos=dim2, y_pos=dim3)
	obs_norm = normalizer.normalizers["observations"]
	self.register_buffer("obs_mins", torch.tensor(obs_norm.mins, dtype=torch.float32))
	self.register_buffer("obs_maxs", torch.tensor(obs_norm.maxs, dtype=torch.float32))

	# dynamics model (A, B, c) as buffers
	if dynamics and dynamics_model is not None:
	self.register_buffer("A", torch.tensor(dynamics_model["A"], dtype=torch.float32))
	self.register_buffer("B", torch.tensor(dynamics_model["B"], dtype=torch.float32))
	self.register_buffer("c", torch.tensor(dynamics_model["c"], dtype=torch.float32).squeeze())

	# precompute obstacle tensors
	if constraints:
	circular = PILLARS + NOVEL_CIRCULAR
	centers = torch.tensor([obs.center for obs in circular], dtype=torch.float32)
	radii = torch.tensor([obs.radius for obs in circular], dtype=torch.float32)
	self.register_buffer("obs_centers", centers)
	self.register_buffer("obs_radii", radii)

	if NOVEL_PLANAR:
	slopes = torch.tensor([obs.slope for obs in NOVEL_PLANAR], dtype=torch.float32)
	intercepts = torch.tensor([obs.intercept for obs in NOVEL_PLANAR], dtype=torch.float32)
	self.register_buffer("planar_slopes", slopes)
	self.register_buffer("planar_intercepts", intercepts)
	else:
	self.planar_slopes = None

	def _unnorm_obs(self, obs_norm):
	"""Unnormalize observations from [-1,1] to real-world coords.

	Args:
	obs_norm: (..., state_dim) normalized observations
	Returns:
	(..., state_dim) real-world observations
	"""
	return (obs_norm + 1.0) / 2.0 * (self.obs_maxs - self.obs_mins) + self.obs_mins

	def forward(self, x):
	"""Compute value for a batch of trajectories.

	Args:
	x: (B, H, action_dim + state_dim) normalized trajectories

	Returns:
	value: (B,) scalar value per trajectory (higher is better)
	"""
	B = x.shape[0]
	actions = x[:, :, :self.action_dim] # (B, H, action_dim)
	obs_norm = x[:, :, self.action_dim:] # (B, H, state_dim)
	obs_real = self._unnorm_obs(obs_norm) # (B, H, state_dim)

	value = torch.zeros(B, device=x.device)

	# --- objective ---
	if self.objective == "distance":
	# terminal y-position cost: -(y_terminal - y_goal)^2
	y_terminal = obs_real[:, -1, 3] # (B,) — y_pos is dim 3
	value = value + self.objective_weight * (-(y_terminal - self.y_goal) ** 2)

	# --- obstacle penalty ---
	if self.constraints and self.obstacle_weight > 0:
	value = value - self.obstacle_weight * self._obstacle_penalty(obs_real)

	# --- dynamics penalty ---
	if self.dynamics and self.dynamics_weight > 0:
	value = value - self.dynamics_weight * self._dynamics_penalty(actions, obs_norm)

	return value

	def _obstacle_penalty(self, obs_real):
	"""Compute squared obstacle violation penalty.

	Args:
	obs_real: (B, H, state_dim) real-world observations
	Returns:
	penalty: (B,) sum of squared violations
	"""
	x_pos = obs_real[:, :, 2] # (B, H)
	y_pos = obs_real[:, :, 3] # (B, H)

	penalty = torch.zeros(obs_real.shape[0], device=obs_real.device)

	# circular obstacles: violation = clamp(r + margin - dist, min=0)
	# obs_centers: (N_obs, 2), obs_radii: (N_obs,)
	for i in range(self.obs_centers.shape[0]):
	cx, cy = self.obs_centers[i]
	r = self.obs_radii[i]
	dist = torch.sqrt((x_pos - cx) 2 + (y_pos - cy) 2) # (B, H)
	violation = torch.clamp(r + self.obstacle_margin - dist, min=0) # (B, H)
	penalty = penalty + (violation ** 2).sum(dim=1) # (B,)

	# planar obstacles: violation = clamp(-(slope * x + intercept - y) + margin, min=0)
	# feasible: slope * x + intercept - y >= 0
	if self.planar_slopes is not None:
	for i in range(self.planar_slopes.shape[0]):
	s = self.planar_slopes[i]
	b = self.planar_intercepts[i]
	feasibility = s * x_pos + b - y_pos # (B, H) — positive = feasible
	violation = torch.clamp(self.obstacle_margin - feasibility, min=0) # (B, H)
	penalty = penalty + (violation ** 2).sum(dim=1)

	return penalty

	def _dynamics_penalty(self, actions, obs_norm):
	"""Compute squared dynamics residual penalty in normalized space.

	s_{t+1} = A @ s_t + B @ a_t + c (fitted linear model)

	Args:
	actions: (B, H, action_dim) normalized actions
	obs_norm: (B, H, state_dim) normalized observations
	Returns:
	penalty: (B,) sum of squared residuals
	"""
	# s_t for t=0..H-2, s_{t+1} for t=1..H-1
	s_t = obs_norm[:, :-1, :] # (B, H-1, state_dim)
	s_next = obs_norm[:, 1:, :] # (B, H-1, state_dim)
	a_t = actions[:, :-1, :] # (B, H-1, action_dim)

	# predicted: A @ s_t + B @ a_t + c
	# A: (state_dim, state_dim), B: (state_dim, action_dim), c: (state_dim,)
	pred = (torch.einsum("ij,bhj->bhi", self.A, s_t)
	+ torch.einsum("ij,bhj->bhi", self.B, a_t)
	+ self.c) # (B, H-1, state_dim)

	residual = pred - s_next # (B, H-1, state_dim)
	penalty = (residual ** 2).sum(dim=(1, 2)) # (B,)
	return penalty