75 lines
1.7 KiB
Python
75 lines
1.7 KiB
Python
"""
|
|
Utility functions for state representation and environment interaction.
|
|
"""
|
|
|
|
import torch
|
|
from typing import List, Tuple, Any
|
|
|
|
|
|
def encode_state(state: Any) -> torch.Tensor:
|
|
"""
|
|
Encode a generic state into a torch tensor.
|
|
|
|
Parameters
|
|
----------
|
|
state : Any
|
|
The state to encode. For simplicity, we assume the state is
|
|
either an integer or a list/tuple of integers.
|
|
|
|
Returns
|
|
-------
|
|
torch.Tensor
|
|
A 1-D tensor representing the state.
|
|
"""
|
|
if isinstance(state, int):
|
|
return torch.tensor([state], dtype=torch.float32)
|
|
elif isinstance(state, (list, tuple)):
|
|
return torch.tensor(state, dtype=torch.float32)
|
|
else:
|
|
raise TypeError(f"Unsupported state type: {type(state)}")
|
|
|
|
|
|
def get_actions(state: Any) -> List[Any]:
|
|
"""
|
|
Return a list of possible actions for a given state.
|
|
|
|
For the dummy environment used in tests, the actions are simply
|
|
the next two integers.
|
|
|
|
Parameters
|
|
----------
|
|
state : Any
|
|
Current state.
|
|
|
|
Returns
|
|
-------
|
|
List[Any]
|
|
List of possible actions.
|
|
"""
|
|
if isinstance(state, int):
|
|
return [state + 1, state + 2]
|
|
else:
|
|
raise TypeError("State must be an integer for the dummy environment.")
|
|
|
|
|
|
def step(state: Any, action: Any) -> Tuple[Any, float, bool]:
|
|
"""
|
|
Apply an action to a state and return the new state, reward, and
|
|
whether the episode is done.
|
|
|
|
Parameters
|
|
----------
|
|
state : Any
|
|
Current state.
|
|
action : Any
|
|
Action to apply.
|
|
|
|
Returns
|
|
-------
|
|
Tuple[Any, float, bool]
|
|
New state, reward, done flag.
|
|
"""
|
|
new_state = action
|
|
reward = 1.0 if new_state == 10 else 0.0
|
|
done = new_state >= 10
|
|
return new_state, reward, done |