Spaces:
Running on Zero
Running on Zero
| """ | |
| Random walk noise augmentation for physics simulation. | |
| Based on Learning to Simulate (Sanchez-Gonzalez et al., 2020) | |
| """ | |
| import torch | |
| import numpy as np | |
| from typing import Optional, Tuple, Dict | |
| def get_random_walk_noise_for_position_sequence( | |
| position_sequence: torch.Tensor, | |
| noise_std_last_step: float, | |
| dt: float = 1.0 | |
| ) -> torch.Tensor: | |
| """ | |
| Generate random walk noise for position sequences. | |
| This noise mimics the accumulation of errors during rollout, | |
| making the model more robust to its own prediction errors. | |
| Args: | |
| position_sequence: Tensor of shape (time, num_particles, 3) or (batch, time, num_particles, 3) | |
| noise_std_last_step: Standard deviation of noise at the last step | |
| dt: Time step for velocity computation | |
| Returns: | |
| position_noise: Tensor of same shape as position_sequence | |
| """ | |
| # Handle both 3D and 4D tensors | |
| if position_sequence.dim() == 3: | |
| # Add batch dimension for consistency | |
| position_sequence = position_sequence.unsqueeze(0) | |
| squeeze_batch = True | |
| else: | |
| squeeze_batch = False | |
| batch_size, seq_len, num_particles, dim = position_sequence.shape | |
| # We want the noise scale in the velocity at the last step to be fixed. | |
| # Because we are going to compose noise at each step using a random walk: | |
| # std_last_step^2 = num_velocities * std_each_step^2 | |
| # so to keep std_last_step fixed, we apply at each step: | |
| # std_each_step = std_last_step / sqrt(num_velocities) | |
| num_velocities = seq_len - 1 # Number of velocity steps | |
| if num_velocities <= 0: | |
| return torch.zeros_like(position_sequence) | |
| # Generate independent noise for each velocity step | |
| velocity_shape = (batch_size, num_velocities, num_particles, dim) | |
| std_each_step = noise_std_last_step / np.sqrt(num_velocities) | |
| # Generate velocity noise | |
| velocity_noise = torch.randn( | |
| velocity_shape, | |
| device=position_sequence.device, | |
| dtype=position_sequence.dtype | |
| ) * std_each_step | |
| # Apply random walk by cumulative sum | |
| velocity_noise_accumulated = torch.cumsum(velocity_noise, dim=1) | |
| # Integrate velocity noise to get position noise | |
| # First position has no noise, then cumsum of velocities | |
| position_noise = torch.zeros_like(position_sequence) | |
| if num_velocities > 0: | |
| # Cumulative sum of accumulated velocities gives position noise | |
| position_noise_from_vel = torch.cumsum(velocity_noise_accumulated * dt, dim=1) | |
| position_noise[:, 1:] = position_noise_from_vel | |
| # Remove batch dimension if we added it | |
| if squeeze_batch: | |
| position_noise = position_noise.squeeze(0) | |
| return position_noise | |
| def apply_random_walk_noise_to_batch( | |
| batch: Dict[str, torch.Tensor], | |
| noise_config: Dict, | |
| is_training: bool = True | |
| ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: | |
| """ | |
| Apply random walk noise to a batch of data. | |
| Args: | |
| batch: Dictionary containing 'body1_positions', 'body2_positions', etc. | |
| noise_config: Configuration dictionary with noise parameters | |
| is_training: Whether in training mode (no noise during validation) | |
| Returns: | |
| noisy_batch: Batch with noise applied to positions | |
| noise_info: Dictionary containing the applied noise for target correction | |
| """ | |
| if not is_training or not noise_config.get('random_walk', {}).get('enable', False): | |
| return batch, {} | |
| noise_std = noise_config['random_walk'].get('noise_std', 0.0001) | |
| dt = noise_config.get('dt', 1.0) | |
| # Copy batch to avoid modifying original | |
| noisy_batch = {k: v.clone() if torch.is_tensor(v) else v for k, v in batch.items()} | |
| noise_info = {} | |
| # Apply noise to body1 and body2 positions | |
| for body_key in ['body1_positions', 'body2_positions']: | |
| if body_key in batch: | |
| positions = batch[body_key] | |
| # Generate random walk noise | |
| position_noise = get_random_walk_noise_for_position_sequence( | |
| positions, noise_std, dt | |
| ) | |
| # Apply noise to positions | |
| noisy_batch[body_key] = positions + position_noise | |
| # Store noise for target correction | |
| noise_key = body_key.replace('positions', 'noise') | |
| noise_info[noise_key] = position_noise | |
| # Recompute velocities from noisy positions | |
| vel_key = body_key.replace('positions', 'velocities') | |
| if vel_key in batch: | |
| # Compute velocities from noisy positions | |
| noisy_positions = noisy_batch[body_key] | |
| # Handle both 3D (seq_len, num_particles, 3) and 4D tensors | |
| if noisy_positions.dim() == 3: | |
| seq_len = noisy_positions.shape[0] | |
| if seq_len > 1: | |
| # Finite differences between consecutive noisy positions | |
| velocities = torch.zeros_like(batch[vel_key]) | |
| vel_seq_len = velocities.shape[0] | |
| if vel_seq_len > 0: | |
| # First velocity from first two positions | |
| velocities[0] = (noisy_positions[1] - noisy_positions[0]) / dt | |
| # Middle velocities | |
| for i in range(1, min(vel_seq_len-1, seq_len-2)): | |
| velocities[i] = (noisy_positions[i+1] - noisy_positions[i]) / dt | |
| # Last velocity if needed | |
| if vel_seq_len > 1 and seq_len > 2: | |
| velocities[-1] = (noisy_positions[-1] - noisy_positions[-2]) / dt | |
| noisy_batch[vel_key] = velocities | |
| else: | |
| # 4D case (batch dimension) | |
| if noisy_positions.shape[1] > 1: | |
| velocities = torch.zeros_like(batch[vel_key]) | |
| if velocities.shape[1] > 0: | |
| velocities[:, 0] = (noisy_positions[:, 1] - noisy_positions[:, 0]) / dt | |
| for i in range(1, min(velocities.shape[1]-1, noisy_positions.shape[1]-2)): | |
| velocities[:, i] = (noisy_positions[:, i+1] - noisy_positions[:, i]) / dt | |
| if velocities.shape[1] > 1 and noisy_positions.shape[1] > 2: | |
| velocities[:, -1] = (noisy_positions[:, -1] - noisy_positions[:, -2]) / dt | |
| noisy_batch[vel_key] = velocities | |
| # Recompute accelerations from noisy positions | |
| accel_key = body_key.replace('positions', 'accelerations') | |
| if accel_key in batch: | |
| noisy_positions = noisy_batch[body_key] | |
| # Handle both 3D and 4D tensors | |
| if noisy_positions.dim() == 3: | |
| seq_len = noisy_positions.shape[0] | |
| if seq_len >= 3: | |
| # Compute accelerations using second-order differences | |
| accelerations = torch.zeros_like(batch[accel_key]) | |
| accel_seq_len = accelerations.shape[0] | |
| for i in range(accel_seq_len): | |
| if i + 1 < seq_len - 1: | |
| accelerations[i] = ( | |
| noisy_positions[i+2] - 2 * noisy_positions[i+1] + noisy_positions[i] | |
| ) / (dt * dt) | |
| noisy_batch[accel_key] = accelerations | |
| else: | |
| # 4D case (batch dimension) | |
| if noisy_positions.shape[1] >= 3: | |
| accelerations = torch.zeros_like(batch[accel_key]) | |
| for i in range(accelerations.shape[1]): | |
| if i + 1 < noisy_positions.shape[1] - 1: | |
| accelerations[:, i] = ( | |
| noisy_positions[:, i+2] - 2 * noisy_positions[:, i+1] + noisy_positions[:, i] | |
| ) / (dt * dt) | |
| noisy_batch[accel_key] = accelerations | |
| return noisy_batch, noise_info | |
| def adjust_targets_for_noise( | |
| targets: Dict[str, torch.Tensor], | |
| noise_info: Dict[str, torch.Tensor], | |
| gamma: float = 1.0 | |
| ) -> Dict[str, torch.Tensor]: | |
| """ | |
| Adjust targets to account for input noise, so the model learns to correct it. | |
| Args: | |
| targets: Dictionary containing target positions/accelerations | |
| noise_info: Dictionary containing the noise applied to inputs | |
| gamma: Correction factor (1.0 = full correction, 0.0 = no correction) | |
| Returns: | |
| adjusted_targets: Targets adjusted for noise correction | |
| """ | |
| adjusted_targets = {k: v.clone() if torch.is_tensor(v) else v for k, v in targets.items()} | |
| # Adjust position targets | |
| for body in ['body1', 'body2']: | |
| noise_key = f'{body}_noise' | |
| pos_key = f'{body}_positions' | |
| if noise_key in noise_info and pos_key in targets: | |
| # The target should include the noise at the last input step | |
| # so the model learns to predict outputs that correct for the noise | |
| last_noise = noise_info[noise_key][:, -1] if noise_info[noise_key].shape[1] > 0 else 0 | |
| adjusted_targets[pos_key] = targets[pos_key] + gamma * last_noise | |
| return adjusted_targets | |
| def apply_independent_noise( | |
| batch: Dict[str, torch.Tensor], | |
| noise_config: Dict, | |
| is_training: bool = True | |
| ) -> Dict[str, torch.Tensor]: | |
| """ | |
| Apply independent Gaussian noise to velocities and other features. | |
| Independent noise model. | |
| Args: | |
| batch: Input batch | |
| noise_config: Noise configuration | |
| is_training: Whether in training mode | |
| Returns: | |
| Noisy batch | |
| """ | |
| if not is_training or not noise_config.get('independent', {}).get('enable', False): | |
| return batch | |
| noisy_batch = {k: v.clone() if torch.is_tensor(v) else v for k, v in batch.items()} | |
| # Apply velocity noise | |
| vel_noise_std = noise_config['independent'].get('velocity_noise', 0.0) | |
| if vel_noise_std > 0: | |
| for key in ['body1_velocities', 'body2_velocities']: | |
| if key in batch: | |
| noise = torch.randn_like(batch[key]) * vel_noise_std | |
| noisy_batch[key] = batch[key] + noise | |
| return noisy_batch |