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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164 | from dataclasses import dataclass, field
from typing import Dict
import numpy as np
import torch
from incar.extensions.ai import BasePolicy, LRSchedulerConfig, OptimizerConfig, PolicyConfig, NormalizationMode, action_tensor_to_dict
from incar.common import FeatureType, ProcessHook
@PolicyConfig.register_subclass("my_custom_policy")
@dataclass
class MyCustomPolicyConfig(PolicyConfig):
# You can add fields for your policy configuration here. The base config already defines some
# fields, like `dt`, `n_obs_steps` and `n_action_steps`. You can override standard values here,
# or during training, values can be overriden with flags, e.g. --policy.device = "cpu"
dt: float = 0.1
n_obs_steps: int = 1
n_action_steps: int = 10
# default normalization modes used for features if they do not explicitly provide their own
# normalization mode
normalization_mapping: dict[FeatureType, NormalizationMode] = field(
default_factory=lambda: {
FeatureType.VISUAL: NormalizationMode.VIDEO_ZERO_ONE,
FeatureType.STATE: NormalizationMode.MIN_MAX,
FeatureType.ACTION: NormalizationMode.MIN_MAX,
}
)
@property
def observation_relative_indices(self) -> list:
raise NotImplementedError
@property
def action_relative_indices(self) -> list:
raise NotImplementedError
"""
Here you can configure a default optimizer for your policy, which is used when no optimizer
is configured in the training config
"""
def get_default_optimizer(self) -> OptimizerConfig:
raise NotImplementedError
"""
Here you can configure a default scheduler for your policy, which is used when no scheduler
is configured in the training config
"""
def get_default_scheduler(self) -> LRSchedulerConfig | None:
raise NotImplementedError
"""
This can be used to validate that the features specified in the config are compatible with the policy implementation.
For example, if a certain policy architecture requires a fixed number of visual features, this can be checked here
and an error raised if the config is not compatible.
"""
def validate_features(self) -> None:
pass
"""
Used when certain parameters are different during inference from training. For example, if during
training a random crop is used but during inference you want to use a center crop, you can set
this here.
"""
def set_inference_params(self) -> None:
pass
def build_policy(self, stats):
return MyCustomPolicy(self, stats)
def build_policy_from_existing_model(self, model_path):
return MyCustomPolicy.load_from_safetensor(
MyCustomPolicy(self),
model_path,
self.device
)
"""
This is an example to show the required interface for your policy. The implementation can highly differ
based on your method and implementation. For clear examples, please check the documentation and
baselines provided in the incar_baselines package: https://github.com/INCAR-Robotics/incar_baselines
"""
class MyCustomPolicy(BasePolicy):
def __init__(
self,
config: MyCustomPolicyConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
):
super().__init__(config, dataset_stats)
# Initialise your policy, observation queues, etc
def get_optim_params(self) -> dict:
# This function has to return a dict containing the parameters that can be optimized.
# e.g. self.model.parameters() if this class contains a model that derives from `torch.nn.Module`
raise NotImplementedError
def queue_observations(self, frame: dict[str, torch.Tensor]):
# These are incoming observations. They should be stored so that `perform_inference` can access them
# The frame consists of {key: value} pairs containing the last received observation for each feature.
# You can see which feature keys you are receiving in the GUI. The shape of `value` is equal to the
# shape of the feature, e.g for video features the shape is (C,H,W).
raise NotImplementedError
@torch.no_grad
def perform_inference(self) -> torch.Tensor:
# This function has to return the predicted action dict. How this is done can vary, but here are some
# likely steps:
# Make a shallow copy of your observations first, since otherwise the reference can be overwritten by
# other incoming frames
# Each frame in the observation queue needs to be processed. A single frame can be processed like so:
self.config.preprocessing.process(obs, ProcessHook.OBSERVATION)
# We process frames here instead of in `queue_observations`, since `queue_observations` runs in the main
# thread - this could cause issues when processing steps include expensive operations such as SAM,
# AnyDepth or Keypoint extraction
# Add a batch dimension and send to the correct device
for key in obs:
obs[key] = obs[key].unsqueeze(0).type(torch.float32).to(self.config.device)
# Normalize the inputs
obs = self.normalize_inputs(obs)
# Concatenate features over a n_obs dimension if multiple observations are used
actions = ... # Get actions with inference
# Get the actions back to dict where they are keyed by your features
action_dict = action_tensor_to_dict(actions[0], self.config.action_features)
# Unnormalize output
action_dict = self.unnormalize_outputs(action_dict)
# Send back to CPU and ensure correct shape
for key, value in action_dict.items():
action_dict[key] = value.to("cpu").numpy().squeeze()
if action_dict[key].ndim == 1:
action_dict[key] = np.expand_dims(action_dict[key], -1)
return action_dict
@torch.no_grad
def validate_batch(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
# The frame consists of {key: value} pairs for each input feature of the policy.
# The shape of `value` is (Batch, n_obs, data_shape), e.g for video features the shape is (Batch,n_obs,C,H,W).
batch = dict(batch) # Make copy
batch = self.normalize_inputs(batch)
batch = self.normalize_targets(batch)
action_mse = ...
return action_mse
def forward(self, batch: Dict[str, torch.Tensor]) -> tuple[torch.Tensor, dict | None]:
# The frame consists of {key: value} pairs for each input feature of the policy.
# The shape of `value` is (Batch, n_obs, data_shape), e.g for video features the shape is (Batch,n_obs,C,H,W).
batch = dict(batch) # Make copy
batch = self.normalize_inputs(batch)
batch = self.normalize_targets(batch)
loss = ...
return loss, None
|