File size: 3,628 Bytes
c2dd7f1 7838c93 c2dd7f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | import json
import re
from pathlib import Path
import torch
import torch.nn as nn
from torch import Tensor
from jaxtyping import Float
from huggingface_hub import PyTorchModelHubMixin
class Model(nn.Module, PyTorchModelHubMixin, repo_url="https://huggingface.co/TimWalter/RAM/new/main", paper_url="https://huggingface.co/papers/2606.09108",license="mit",):
"""
RAM model that predicts reachability from a modified Denavit-Hartenberg morphology parametrisation and a pose.
"""
@classmethod
def from_id(cls, model_id: int):
"""
Instantiate a model from its wandb ID.
Args:
model_id: Wandb ID of the model.
Returns:
model: Instantiated model.
"""
model_dir = Path(__file__).parent.parent / "trained_models"
pattern = rf"{model_id}-[a-z]+-[a-z]+"
folder = next((f for f in model_dir.iterdir() if re.match(pattern, f.name)), None)
metadata_path = model_dir / folder / 'metadata.json'
metadata = json.load(open(metadata_path, 'r'))
model = cls(**metadata["hyperparameter"])
model_folder = Path(str(model_dir / folder))
prime = model_folder / "model.pth"
if prime.exists():
old_state_dict =torch.load(prime)
else:
old_state_dict =torch.load(model_folder / "checkpoint.pth")
new_state_dict = {}
# 2. Map legacy state dict keys to your new flattened structure
for key, value in old_state_dict.items():
new_key = key
if key.startswith("encoder.lstm."):
new_key = key.replace("encoder.lstm.", "encoder.")
elif key.startswith("decoder.model."):
new_key = key.replace("decoder.model.", "decoder.")
new_state_dict[new_key] = value
# Load the corrected state dict into the model
model.load_state_dict(new_state_dict)
return model
def __init__(self,
dim_encoding: int = 128,
num_encoder_layers: int = 1,
drop_prob: float = 0.0,
dim_decoder: int = 1792,
num_decoder_layer: int = 8):
"""
Initialise the model.
Args:
dim_encoding: The dimension of the latent morphology encoding.
num_encoder_layers: Number of LSTM layers.
drop_prob: Dropout probability of the LSTM.
dim_decoder: Hidden dimension of the MLP.
num_decoder_layer: Number of layers of the MLP.
"""
super().__init__()
self.encoder = nn.LSTM(3, dim_encoding, num_encoder_layers, dropout=drop_prob, batch_first=True, bias=False)
self.decoder = nn.Sequential(
nn.Linear(9 + dim_encoding, dim_decoder),
nn.ReLU(),
*[nn.Sequential(nn.Linear(dim_decoder, dim_decoder), nn.ReLU())
for _ in range(num_decoder_layer)],
nn.Linear(dim_decoder, 1)
)
def forward(self, morph: Float[Tensor, "batch seq 3"], pose: Float[Tensor, "batch 9"]) -> Float[Tensor, "batch"]:
"""
Predict reachability.
Args:
morph: Morphology description.
pose: Pose as vector encoded.
Returns:
Reachability logit.
"""
latent = self.encoder(morph)[1][0][-1]
logit = self.decoder(torch.cat([pose, latent], dim=-1)).squeeze(-1)
return logit
@torch.inference_mode()
def predict(self, *args, **kwargs):
"""
forward with inference mode
"""
return self.forward(*args, **kwargs)
|