def state_actions_to_input_data(self, state, action):
"""Process state-action pairs."""
if self.discrete_state:
state = one_hot_encode(state, num_classes=self.num_states)
if self.discrete_action:
action = one_hot_encode(action, num_classes=self.num_actions)
if self.input_transform is not None:
state = self.input_transform(state)
state_action = torch.cat((state, action), dim=-1)
return state_action
def forward(self, state, action=torch.tensor(float("nan"))):
"""Get value of the value-function at a given state."""
if self.input_transform is not None:
state = self.input_transform(state)
if self.discrete_state:
state = one_hot_encode(state, self.num_states)
return self.nn(state).squeeze(-1)
def set_value(self, state, new_value):
"""Set value to value function at a given state."""
if new_value.ndim < 1 or new_value.shape[
-1] != self.nn.head.weight.shape[0]:
new_value = torch.log(
one_hot_encode(new_value, num_classes=self.num_actions) +
1e-12)
with torch.no_grad():
self.nn.head.weight[:, state] = new_value
def _preprocess_state(self, state):
"""Pre-process state before input to neural network."""
if self.input_transform is not None: # Apply input transform.
state = self.input_transform(state)
if self.discrete_state: # One hot encode discrete states.
state = one_hot_encode(state, num_classes=self.num_states)
if self.goal is not None: # concatenate goal to state.
goal = self.goal.repeat(*state.shape[:-1], 1)
state = torch.cat((state, goal), dim=-1)
return state
def forward(self, state, action=torch.tensor(float("nan"))):
"""Get value of the value-function at a given state.
Parameters
----------
state: torch.Tensor
action: torch.Tensor
Returns
-------
value: torch.Tensor
"""
if self.discrete_state:
state = one_hot_encode(state, self.num_states)
if self.input_transform is not None:
state = self.input_transform(state)
if torch.isnan(action).all():
if not self.discrete_action:
raise NotImplementedError
action_value = self.nn(state)
return action_value
if self.discrete_action:
action = action.unsqueeze(-1).long()
if action.dim() < state.dim():
resqueeze = True
action = action.unsqueeze(0)
else:
resqueeze = False
if not self.discrete_action:
state_action = torch.cat((state, action), dim=-1)
return self.nn(state_action).squeeze(-1)
else:
out = self.nn(state).gather(-1, action).squeeze(-1)
if resqueeze:
return out.squeeze(0)
else:
return out
def embeddings(self, state):
"""Get embeddings of the value-function at a given state."""
if self.discrete_state:
state = one_hot_encode(state, self.num_states)
return self.nn.last_layer_embeddings(state).squeeze(-1)