Skip to content

Huggingface integration #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Huggingface integration #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
609862b
hf integration
spirosperos Aug 8, 2025
467b54b
wandb integration
spirosperos Aug 8, 2025
ef9ccc1
Merge branch 'main' into feature-hf
spirosperos Aug 8, 2025
2ead943
fix
spirosperos Aug 8, 2025
6a76911
Merge branch 'feature-hf' of github.com:SpesRobotics/alignit into fea…
spirosperos Aug 8, 2025
0820937
test fix
spirosperos Aug 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@
*.egg-info
data
alignnet_model.pth

wandb/
23 changes: 23 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,27 @@ pip install -e .

# Run tests
python -m pytest
```

### HuggingFace Hub Integration

```bash
# Record and push to HuggingFace Hub
python -m alignit.record --dataset.hf_username=username --dataset.hf_dataset_name=dataset-name

# Train using HuggingFace dataset
python -m alignit.train --dataset.hf_dataset_name=username/dataset-name --model.path=./data/test_model.pth

# Visualize HuggingFace dataset
python -m alignit.visualize --dataset.hf_dataset_name=username/dataset-name
```

### Weights & Biases Integration

```bash
# Train with wandb experiment tracking
python -m alignit.train --wandb_project=my-robot-project --wandb_run_name=experiment-1

# Combine HuggingFace Hub and wandb
python -m alignit.train --dataset.hf_dataset_name=username/dataset-name --wandb_project=my-robot-project
```
160 changes: 47 additions & 113 deletions alignit/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,158 +8,92 @@
@dataclass
class DatasetConfig:
"""Configuration for dataset paths and loading."""

path: str = field(
default="./data/duck", metadata={"help": "Path to the dataset directory"}
)
path: str = field(default="./data/duck", metadata={"help": "Path to the dataset directory"})
hf_username: Optional[str] = field(default=None, metadata={"help": "Hugging Face username for dataset sync"})
hf_dataset_name: Optional[str] = field(default=None, metadata={"help": "Hugging Face dataset name for sync"})
hf_token: Optional[str] = field(default=None, metadata={"help": "Hugging Face token for authentication"})
prefer_local: bool = field(default=True, metadata={"help": "Prefer local dataset over Hugging Face Hub"})
push_to_hub: bool = field(default=False, metadata={"help": "Automatically push datasets to Hugging Face Hub after saving"})


@dataclass
class ModelConfig:
"""Configuration for AlignNet model."""

backbone: str = field(
default="efficientnet_b0",
metadata={"help": "Backbone architecture: 'efficientnet_b0' or 'resnet18'"},
)
backbone_weights: str = field(
default="DEFAULT", metadata={"help": "Backbone weights: 'DEFAULT' or None"}
)
use_vector_input: bool = field(
default=False, metadata={"help": "Whether to use vector input"}
)
fc_layers: List[int] = field(
default_factory=lambda: [256, 128],
metadata={"help": "Hidden layer sizes for FC head"},
)
vector_hidden_dim: int = field(
default=64, metadata={"help": "Output dimension of vector MLP"}
)
output_dim: int = field(
default=9,
metadata={"help": "Final output dimension (3 translation + 6 rotation)"},
)
feature_agg: str = field(
default="mean", metadata={"help": "Feature aggregation method: 'mean' or 'max'"}
)
path: str = field(
default="alignnet_model.pth",
metadata={"help": "Path to save/load trained model"},
)
backbone: str = field(default="efficientnet_b0", metadata={"help": "Backbone architecture: 'efficientnet_b0' or 'resnet18'"})
backbone_weights: str = field(default="DEFAULT", metadata={"help": "Backbone weights: 'DEFAULT' or None"})
use_vector_input: bool = field(default=False, metadata={"help": "Whether to use vector input"})
fc_layers: List[int] = field(default_factory=lambda: [256, 128], metadata={"help": "Hidden layer sizes for FC head"})
vector_hidden_dim: int = field(default=64, metadata={"help": "Output dimension of vector MLP"})
output_dim: int = field(default=9, metadata={"help": "Final output dimension (3 translation + 6 rotation)"})
feature_agg: str = field(default="mean", metadata={"help": "Feature aggregation method: 'mean' or 'max'"})
path: str = field(default="alignnet_model.pth", metadata={"help": "Path to save/load trained model"})


@dataclass
class TrajectoryConfig:
"""Configuration for spiral trajectory generation."""

z_step: float = field(
default=0.002, metadata={"help": "Z step size for spiral trajectory"}
)
radius_step: float = field(
default=0.001, metadata={"help": "Radius step size for spiral trajectory"}
)
num_steps: int = field(
default=50, metadata={"help": "Number of steps in spiral trajectory"}
)
z_step: float = field(default=0.0007, metadata={"help": "Z step size for spiral trajectory"})
radius_step: float = field(default=0.001, metadata={"help": "Radius step size for spiral trajectory"})
num_steps: int = field(default=100, metadata={"help": "Number of steps in spiral trajectory"})
cone_angle: float = field(default=30.0, metadata={"help": "Cone angle in degrees"})
visible_sweep: float = field(
default=60.0, metadata={"help": "Visible sweep angle in degrees"}
)
viewing_angle_offset: float = field(
default=-120.0, metadata={"help": "Viewing angle offset in degrees"}
)
angular_resolution: float = field(
default=10.0, metadata={"help": "Angular resolution in degrees"}
)
include_cone_poses: bool = field(
default=False, metadata={"help": "Include cone poses in trajectory"}
)
lift_height_before_spiral: float = field(
default=0.01, metadata={"help": "Lift height before spiral in meters"}
)
visible_sweep: float = field(default=60.0, metadata={"help": "Visible sweep angle in degrees"})
viewing_angle_offset: float = field(default=-120.0, metadata={"help": "Viewing angle offset in degrees"})
angular_resolution: float = field(default=10.0, metadata={"help": "Angular resolution in degrees"})
include_cone_poses: bool = field(default=False, metadata={"help": "Include cone poses in trajectory"})
lift_height_before_spiral: float = field(default=0.01, metadata={"help": "Lift height before spiral in meters"})


@dataclass
class RecordConfig:
"""Configuration for data recording."""

dataset: DatasetConfig = field(default_factory=DatasetConfig)
trajectory: TrajectoryConfig = field(default_factory=TrajectoryConfig)
episodes: int = field(default=10, metadata={"help": "Number of episodes to record"})
lin_tol_alignment: float = field(
default=0.015, metadata={"help": "Linear tolerance for alignment servo"}
)
ang_tol_alignment: float = field(
default=0.015, metadata={"help": "Angular tolerance for alignment servo"}
)
lin_tol_trajectory: float = field(
default=0.05, metadata={"help": "Linear tolerance for trajectory servo"}
)
ang_tol_trajectory: float = field(
default=0.05, metadata={"help": "Angular tolerance for trajectory servo"}
)
episodes: int = field(default=20, metadata={"help": "Number of episodes to record"})
lin_tol_alignment: float = field(default=0.015, metadata={"help": "Linear tolerance for alignment servo"})
ang_tol_alignment: float = field(default=0.015, metadata={"help": "Angular tolerance for alignment servo"})
lin_tol_trajectory: float = field(default=0.05, metadata={"help": "Linear tolerance for trajectory servo"})
ang_tol_trajectory: float = field(default=0.05, metadata={"help": "Angular tolerance for trajectory servo"})


@dataclass
class TrainConfig:
"""Configuration for model training."""

dataset: DatasetConfig = field(default_factory=DatasetConfig)
model: ModelConfig = field(default_factory=ModelConfig)
batch_size: int = field(default=8, metadata={"help": "Training batch size"})
learning_rate: float = field(
default=1e-4, metadata={"help": "Learning rate for optimizer"}
)
learning_rate: float = field(default=1e-4, metadata={"help": "Learning rate for optimizer"})
epochs: int = field(default=100, metadata={"help": "Number of training epochs"})
test_size: float = field(
default=0.2, metadata={"help": "Fraction of data for testing"}
test_size: float = field(default=0.2, metadata={"help": "Fraction of data for testing"})
random_seed: int = field(default=42, metadata={"help": "Random seed for train/test split"})
# Weights & Biases integration
wandb_project: Optional[str] = field(
default=None, metadata={"help": "Weights & Biases project name"}
)
wandb_run_name: Optional[str] = field(
default=None, metadata={"help": "Weights & Biases run name"}
)
random_seed: int = field(
default=42, metadata={"help": "Random seed for train/test split"}
wandb_tags: Optional[List[str]] = field(
default=None, metadata={"help": "Weights & Biases tags for the run"}
)


@dataclass
class InferConfig:
"""Configuration for inference/alignment."""

model: ModelConfig = field(default_factory=ModelConfig)
start_pose_xyz: List[float] = field(
default_factory=lambda: [0.33, 0.0, 0.35],
metadata={"help": "Starting pose XYZ coordinates"},
)
start_pose_rpy: List[float] = field(
default_factory=lambda: [np.pi, 0.0, 0.0],
metadata={"help": "Starting pose RPY angles"},
)
lin_tolerance: float = field(
default=2e-3, metadata={"help": "Linear tolerance for convergence (meters)"}
)
ang_tolerance: float = field(
default=2, metadata={"help": "Angular tolerance for convergence (degrees)"}
)
max_iterations: Optional[int] = field(
default=None,
metadata={"help": "Maximum iterations before stopping (None = infinite)"},
)
debug_output: bool = field(
default=True, metadata={"help": "Print debug information during inference"}
)
debouncing_count: int = field(
default=5,
metadata={"help": "Number of iterations within tolerance before stopping"},
)
start_pose_xyz: List[float] = field(default_factory=lambda: [0.33, 0.0, 0.35], metadata={"help": "Starting pose XYZ coordinates"})
start_pose_rpy: List[float] = field(default_factory=lambda: [np.pi, 0.0, 0.0], metadata={"help": "Starting pose RPY angles"})
lin_tolerance: float = field(default=2e-3, metadata={"help": "Linear tolerance for convergence (meters)"})
ang_tolerance: float = field(default=2.0, metadata={"help": "Angular tolerance for convergence (degrees)"})
max_iterations: Optional[int] = field(default=None, metadata={"help": "Maximum iterations before stopping (None = infinite)"})
debug_output: bool = field(default=True, metadata={"help": "Print debug information during inference"})
debouncing_count: int = field(default=5, metadata={"help": "Number of consecutive iterations within tolerance before stopping"})


@dataclass
class VisualizeConfig:
"""Configuration for dataset visualization."""

dataset: DatasetConfig = field(default_factory=DatasetConfig)
share: bool = field(default=False, metadata={"help": "Create a public Gradio link"})
server_name: Optional[str] = field(
default=None, metadata={"help": "Server name for Gradio interface"}
)
server_port: Optional[int] = field(
default=None, metadata={"help": "Server port for Gradio interface"}
)
server_name: Optional[str] = field(default=None, metadata={"help": "Server name for Gradio interface"})
server_port: Optional[int] = field(default=None, metadata={"help": "Server port for Gradio interface"})
18 changes: 17 additions & 1 deletion alignit/record.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@
load_from_disk,
concatenate_datasets,
)

from alignit.robots.xarmsim import XarmSim
from alignit.robots.xarm import Xarm
from alignit.utils.zhou import se3_sixd
from alignit.utils.dataset import push_dataset_to_hub
import draccus
from alignit.config import RecordConfig

Expand Down Expand Up @@ -132,6 +132,22 @@ def main(cfg: RecordConfig):
shutil.rmtree(cfg.dataset.path)
shutil.move(temp_path, cfg.dataset.path)

# Push to HuggingFace Hub if configured
if cfg.dataset.hf_username and cfg.dataset.hf_dataset_name:
print(f"Pushing dataset to HuggingFace Hub: {cfg.dataset.hf_username}/{cfg.dataset.hf_dataset_name}")
try:
repo_id = push_dataset_to_hub(
combined_dataset,
cfg.dataset.hf_username,
cfg.dataset.hf_dataset_name,
cfg.dataset.hf_token,
private=True # Make private by default for safety
)
print(f"Successfully pushed dataset to {repo_id}")
except Exception as e:
print(f"Failed to push dataset to HuggingFace Hub: {e}")
print("Dataset saved locally but not uploaded to Hub")

robot.disconnect()


Expand Down
58 changes: 56 additions & 2 deletions alignit/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,14 @@

from alignit.config import TrainConfig
from alignit.models.alignnet import AlignNet
from alignit.utils.dataset import load_dataset_smart

# Optional wandb import
try:
import wandb
WANDB_AVAILABLE = True
except ImportError:
WANDB_AVAILABLE = False


def collate_fn(batch):
Expand All @@ -22,8 +30,28 @@ def main(cfg: TrainConfig):
"""Train AlignNet model using configuration parameters."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load the dataset from disk
dataset = load_from_disk(cfg.dataset.path)
# Initialize Weights & Biases if configured
use_wandb = WANDB_AVAILABLE and cfg.wandb_project is not None
if use_wandb:
wandb.init(
project=cfg.wandb_project,
name=cfg.wandb_run_name,
tags=cfg.wandb_tags,
config=cfg.__dict__
)
print(f"Initialized Weights & Biases project: {cfg.wandb_project}")
elif cfg.wandb_project and not WANDB_AVAILABLE:
print("Warning: wandb project specified but wandb not installed. Install with: pip install wandb")

# Load the dataset from disk or HuggingFace Hub
if cfg.dataset.hf_dataset_name:
print(f"Loading dataset from HuggingFace Hub: {cfg.dataset.hf_dataset_name}")
dataset_path = cfg.dataset.hf_dataset_name
else:
print(f"Loading dataset from disk: {cfg.dataset.path}")
dataset_path = cfg.dataset.path

dataset = load_dataset_smart(dataset_path)

# Create model using config parameters
net = AlignNet(
Expand Down Expand Up @@ -53,7 +81,14 @@ def main(cfg: TrainConfig):
criterion = MSELoss()
net.train()

# Watch model with wandb if enabled
if use_wandb:
wandb.watch(net, log_freq=100)

for epoch in range(cfg.epochs):
epoch_loss = 0.0
num_batches = 0

for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
images = batch["images"]
actions = batch["action"].to(device)
Expand All @@ -78,13 +113,32 @@ def main(cfg: TrainConfig):
loss = criterion(outputs, actions)
loss.backward()
optimizer.step()

epoch_loss += loss.item()
num_batches += 1

tqdm.write(f"Loss: {loss.item():.4f}")

# Calculate average epoch loss
avg_epoch_loss = epoch_loss / num_batches

# Log to wandb if enabled
if use_wandb:
wandb.log({
"epoch": epoch + 1,
"loss": avg_epoch_loss,
"learning_rate": cfg.learning_rate
})

# Save the trained model
torch.save(net.state_dict(), cfg.model.path)
tqdm.write(f"Model saved as {cfg.model.path}")

print("Training complete.")

# Finish wandb run
if use_wandb:
wandb.finish()


if __name__ == "__main__":
Expand Down
Loading