mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-19 12:57:58 +00:00
326 lines
13 KiB
Python
326 lines
13 KiB
Python
import logging
|
|
import queue
|
|
from concurrent.futures import ThreadPoolExecutor, TimeoutError
|
|
|
|
import numpy as np
|
|
from dotenv import load_dotenv
|
|
from openai import OpenAI
|
|
from scipy.spatial import cKDTree
|
|
from sklearn.decomposition import PCA
|
|
|
|
from atroposlib.envs.base import BaseEnv
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
|
|
load_dotenv()
|
|
|
|
|
|
class StarMapCompressionEnv(BaseEnv):
|
|
def __init__(self, data_path, views_path):
|
|
logging.info("Initializing StarMapCompressionEnv")
|
|
self.original_data = np.load(data_path)
|
|
self.views = np.load(views_path)
|
|
logging.info(
|
|
f"Original data shape: {self.original_data.shape}, Views shape: {self.views.shape}"
|
|
)
|
|
logging.info(
|
|
f"Original data range: min={self.original_data.min(axis=0)}, max={self.original_data.max(axis=0)}"
|
|
)
|
|
|
|
# Sample points based on density
|
|
self.sampled_data = self._density_sample(self.original_data)
|
|
logging.info(f"Sampled data shape: {self.sampled_data.shape}")
|
|
|
|
# Apply PCA to reduce to 2D
|
|
self.pca_data, self.pca_views = self._apply_pca(self.sampled_data, self.views)
|
|
logging.info(
|
|
f"PCA data shape: {self.pca_data.shape}, PCA views shape: {self.pca_views.shape}"
|
|
)
|
|
|
|
# Compute adaptive density threshold
|
|
data_range = self.original_data.max(axis=0) - self.original_data.min(axis=0)
|
|
cell_volume = (
|
|
np.prod(data_range[data_range > 0]) if np.any(data_range > 0) else 1.0
|
|
)
|
|
self.density_threshold = max(len(self.pca_data) / cell_volume * 0.001, 1e-6)
|
|
logging.info(f"Adaptive density threshold: {self.density_threshold}")
|
|
|
|
# Build initial octree
|
|
self.octree_data = self._build_octree(self.pca_data)
|
|
logging.info(f"Octree data shape: {self.octree_data.shape}")
|
|
|
|
# Quantize the octree data
|
|
self.quantized_data = self._quantize_data(self.octree_data)
|
|
logging.info(f"Quantized data shape: {self.quantized_data.shape}")
|
|
|
|
# Map quantized data back to original points
|
|
self.data = self._map_to_original(self.quantized_data)
|
|
logging.info(f"Final mapped data shape: {self.data.shape}")
|
|
|
|
# Update views to use mapped data
|
|
self.views = self._map_to_original(self.pca_views)
|
|
# Optional: Center views around data mean (uncomment if needed)
|
|
# data_mean = self.original_data.mean(axis=0)
|
|
# self.views = self.views - self.views.mean(axis=0) + data_mean
|
|
logging.info(f"Final views shape: {self.views.shape}")
|
|
|
|
# Scale view radius and grid sizes
|
|
valid_range = data_range[data_range > 0]
|
|
self.view_radius = (
|
|
min(valid_range) * 2.0 if len(valid_range) > 0 else 50.0
|
|
) # ~53.4 to cover min view distance
|
|
self.partition_methods = [
|
|
self.view_radius * 0.5,
|
|
self.view_radius,
|
|
self.view_radius * 1.5,
|
|
]
|
|
logging.info(
|
|
f"View radius: {self.view_radius}, Partition methods: {self.partition_methods}"
|
|
)
|
|
|
|
self.max_steps = 50
|
|
self.client = OpenAI(base_url="http://localhost:9001/v1")
|
|
self.current_method = 0
|
|
|
|
def _density_sample(self, data, sample_fraction=0.1, radius=50.0):
|
|
if len(data) == 0:
|
|
return data
|
|
tree = cKDTree(data)
|
|
density = np.array(
|
|
[len(tree.query_ball_point(point, radius)) for point in data]
|
|
)
|
|
density = density / (density.sum() + 1e-10)
|
|
num_samples = max(1, int(len(data) * sample_fraction))
|
|
indices = np.random.choice(
|
|
len(data), size=num_samples, p=density, replace=False
|
|
)
|
|
return data[indices]
|
|
|
|
def _apply_pca(self, data, views):
|
|
if len(data) == 0:
|
|
return np.array([]).reshape(0, 3), np.array([]).reshape(0, 3)
|
|
pca = PCA(n_components=2)
|
|
pca_data = pca.fit_transform(data)
|
|
pca_views = pca.transform(views)
|
|
pca_data_3d = np.pad(pca_data, ((0, 0), (0, 1)), mode="constant")
|
|
pca_views_3d = np.pad(pca_views, ((0, 0), (0, 1)), mode="constant")
|
|
return pca_data_3d, pca_views_3d
|
|
|
|
def _build_octree(self, data, min_points=2, max_depth=5, density_threshold=None):
|
|
if density_threshold is None:
|
|
density_threshold = self.density_threshold
|
|
|
|
if len(data) < min_points or max_depth <= 0:
|
|
if len(data) > 0:
|
|
return np.mean(data, axis=0, keepdims=True)
|
|
return np.array([]).reshape(0, 3)
|
|
|
|
min_coords = data.min(axis=0)
|
|
max_coords = data.max(axis=0)
|
|
center = (min_coords + max_coords) / 2
|
|
extent = (max_coords - min_coords) / 2
|
|
cell_volume = np.prod(extent[extent > 0]) if np.any(extent > 0) else 1.0
|
|
density = len(data) / cell_volume if cell_volume > 0 else 0
|
|
logging.info(
|
|
f"Octree level {5-max_depth}: density={density}, threshold={density_threshold}"
|
|
)
|
|
|
|
# Skip density pruning for root node
|
|
if max_depth < 5 and density < density_threshold:
|
|
return np.array([]).reshape(0, 3)
|
|
|
|
octants = [[] for _ in range(8)]
|
|
for point in data:
|
|
idx = 0
|
|
if point[0] > center[0]:
|
|
idx |= 1
|
|
if point[1] > center[1]:
|
|
idx |= 2
|
|
if point[2] > center[2]:
|
|
idx |= 4
|
|
octants[idx].append(point)
|
|
|
|
reduced_data = []
|
|
for octant in octants:
|
|
if len(octant) > 0:
|
|
octant_data = np.array(octant)
|
|
subtree = self._build_octree(
|
|
octant_data, min_points, max_depth - 1, density_threshold
|
|
)
|
|
if len(subtree) > 0:
|
|
reduced_data.append(subtree)
|
|
|
|
if not reduced_data:
|
|
return np.array([]).reshape(0, 3)
|
|
return np.vstack(reduced_data)
|
|
|
|
def _quantize_data(self, data, bits=8):
|
|
if len(data) == 0:
|
|
return data
|
|
data_min = data.min(axis=0)
|
|
data_max = data.max(axis=0)
|
|
scale = (2**bits - 1) / (data_max - data_min + 1e-10)
|
|
quantized = np.round((data - data_min) * scale).astype(np.uint8)
|
|
dequantized = (quantized / scale) + data_min
|
|
return dequantized
|
|
|
|
def _map_to_original(self, simplified_data):
|
|
if len(simplified_data) == 0:
|
|
return simplified_data
|
|
tree = cKDTree(self.original_data)
|
|
distances, indices = tree.query(simplified_data, k=1)
|
|
return self.original_data[indices]
|
|
|
|
def reset(self):
|
|
self.step_count = 0
|
|
self.current_method = 0
|
|
return self._get_state()
|
|
|
|
def step(self, action):
|
|
self.step_count += 1
|
|
self.current_method = action
|
|
avg_data_size = self._evaluate_partition()
|
|
# Update data based on grid size
|
|
self.data = self._recompress_data(self.partition_methods[action])
|
|
# Reward balances size, retention, points in view, and quality
|
|
total_points = sum(
|
|
np.sum(np.sqrt(np.sum((self.data - v) ** 2, axis=1)) < self.view_radius)
|
|
for v in self.views
|
|
)
|
|
quality = (
|
|
np.mean([np.min(np.sum((self.data - v) ** 2, axis=1)) for v in self.views])
|
|
if len(self.data) > 0
|
|
else 0
|
|
)
|
|
reward = (
|
|
-avg_data_size / 1000
|
|
+ 5 * len(self.data) / len(self.original_data)
|
|
+ total_points / len(self.original_data)
|
|
- quality / 1e6
|
|
)
|
|
done = self.step_count >= self.max_steps
|
|
logging.info(
|
|
f"Step reward: data_size={avg_data_size}, points={len(self.data)}, "
|
|
f"total_points={total_points}, quality={quality}, reward={reward}"
|
|
)
|
|
return self._get_state(), reward, done, {}
|
|
|
|
def _evaluate_partition(self):
|
|
cell_size = self.partition_methods[self.current_method]
|
|
total_size = 0
|
|
total_points = 0
|
|
for view in self.views:
|
|
distances = np.sqrt(np.sum((self.data - view) ** 2, axis=1))
|
|
points_in_view = np.sum(distances < self.view_radius)
|
|
total_points += points_in_view
|
|
data_size = points_in_view * 32
|
|
cells_per_axis = int(np.ceil(2 * self.view_radius / max(cell_size, 1e-6)))
|
|
num_cells = cells_per_axis**3
|
|
cell_overhead = num_cells * 10
|
|
total_size += data_size + cell_overhead
|
|
logging.info(
|
|
f"View: points_in_view={points_in_view}, num_cells={num_cells}, "
|
|
f"cell_size={cell_size}, data_size={data_size}, cell_overhead={cell_overhead}"
|
|
)
|
|
|
|
avg_data_size = total_size / len(self.views) if len(self.views) > 0 else 0
|
|
logging.info(
|
|
f"Avg data size: {avg_data_size} bytes for grid size {cell_size}, Total points in views: {total_points}"
|
|
)
|
|
return avg_data_size
|
|
|
|
def _recompress_data(self, grid_size):
|
|
# Adjust octree max_depth, min_points, quantization bits, and sampling
|
|
scale = self.view_radius / max(grid_size, 1e-6)
|
|
density_threshold = self.density_threshold / scale
|
|
max_depth = int(3 + 2 * scale) # 3 to 5
|
|
min_points = max(1, int(2 / scale)) # 2 for large grids, higher for small
|
|
bits = max(4, int(6 + 2 * scale)) # 6 to 8 bits
|
|
sample_fraction = min(1.0, 0.3 + scale * 0.3) # 0.3 to 0.6
|
|
logging.info(
|
|
f"Recompress: grid_size={grid_size}, density_threshold={density_threshold}, "
|
|
f"max_depth={max_depth}, min_points={min_points}, bits={bits}, sample_fraction={sample_fraction}"
|
|
)
|
|
octree_data = self._build_octree(
|
|
self.pca_data,
|
|
min_points=min_points,
|
|
max_depth=max_depth,
|
|
density_threshold=density_threshold,
|
|
)
|
|
quantized_data = self._quantize_data(octree_data, bits=bits)
|
|
compressed_data = self._map_to_original(quantized_data)
|
|
# Sample points to introduce variation
|
|
if len(compressed_data) > 0:
|
|
num_samples = max(1, int(len(compressed_data) * sample_fraction))
|
|
indices = np.random.choice(
|
|
len(compressed_data), size=num_samples, replace=False
|
|
)
|
|
compressed_data = compressed_data[indices]
|
|
logging.info(
|
|
f"Recompressed data shape for grid size {grid_size}: {compressed_data.shape}"
|
|
)
|
|
return compressed_data
|
|
|
|
def _get_state(self):
|
|
return {"method": self.current_method, "data_size": len(self.data)}
|
|
|
|
def evaluate(self):
|
|
avg_data_size = self._evaluate_partition()
|
|
return {"avg_data_size": avg_data_size / 1000}
|
|
|
|
def get_next_item(self):
|
|
return self.data
|
|
|
|
def run_rl_step(self, timeout_seconds=60):
|
|
self.current_method = 0
|
|
initial_result = self.evaluate()
|
|
initial_avg_data_size = initial_result["avg_data_size"]
|
|
logging.info(f"Before RL Step: avg_data_size = {initial_avg_data_size} KB")
|
|
|
|
best_reward = float("-inf")
|
|
best_action = np.random.choice(3)
|
|
result_queue = queue.Queue()
|
|
rewards = []
|
|
|
|
def step_with_action(action):
|
|
try:
|
|
state, reward, done, info = self.step(action)
|
|
result_queue.put((reward, action))
|
|
logging.info(
|
|
f"Action {action} (grid size {self.partition_methods[action]}): reward={reward}"
|
|
)
|
|
except Exception as e:
|
|
logging.error(f"Error in action {action}: {e}")
|
|
result_queue.put((float("-inf"), action))
|
|
|
|
with ThreadPoolExecutor(max_workers=3) as executor:
|
|
futures = [executor.submit(step_with_action, action) for action in range(3)]
|
|
for action, future in enumerate(futures):
|
|
try:
|
|
future.result(timeout=timeout_seconds)
|
|
reward, chosen_action = result_queue.get()
|
|
rewards.append((reward, chosen_action))
|
|
except TimeoutError:
|
|
logging.warning(
|
|
f"Action {action} timed out after {timeout_seconds} seconds"
|
|
)
|
|
rewards.append((float("-inf"), action))
|
|
|
|
for reward, action in rewards:
|
|
logging.info(f"Action {action} reward: {reward}")
|
|
|
|
for reward, action in rewards:
|
|
if reward > best_reward or (
|
|
reward == best_reward and np.random.random() < 0.5
|
|
): # Random tiebreaker
|
|
best_reward = reward
|
|
best_action = action
|
|
|
|
self.step(best_action)
|
|
result = self.evaluate()
|
|
logging.info(
|
|
f"After RL Step: Chose grid size {self.partition_methods[best_action]}, "
|
|
f"avg_data_size = {result['avg_data_size']} KB, Reward: {best_reward}"
|
|
)
|