Trajectory Recorder

The trajectory_recorder module provides functionality to save and manage game trajectories, which capture the full sequence of states, actions, and rewards during an episode.

netsecgame.utils.trajectory_recorder

TrajectoryRecorder

Manages the recording and storage of agent trajectories.

Initializes the TrajectoryRecorder.

Parameters:

Name	Type	Description	Default
agent_name	str	The name of the agent.	required
agent_role	str	The role of the agent.	required

Source code in netsecgame/utils/trajectory_recorder.py

def __init__(self, agent_name: str, agent_role: str) -> None:
    """
    Initializes the TrajectoryRecorder.

    Args:
        agent_name (str): The name of the agent.
        agent_role (str): The role of the agent.
    """
    self.agent_name = agent_name
    self.agent_role = agent_role
    self.logger = logging.getLogger(f"TrajectoryRecorder-{agent_name}")
    self._data: Dict[str, Any] = {}
    self.reset()

add_initial_state

Adds the initial state to the trajectory history.

Parameters:

Name	Type	Description	Default
state	GameState	The initial game state.	required

Source code in netsecgame/utils/trajectory_recorder.py

def add_initial_state(self, state: GameState) -> None:
    """
    Adds the initial state to the trajectory history.

    Args:
        state (GameState): The initial game state.
    """
    self._data["trajectory"]["states"].append(state.as_dict)

add_step

Adds a single step to the trajectory.

Parameters:

Name	Type	Description	Default
action	Action	The action taken.	required
reward	float	The reward received.	required
next_state	GameState	The resulting state.	required
end_reason	Optional[str]	Reason for episode end, if applicable.	None

Source code in netsecgame/utils/trajectory_recorder.py

def add_step(self, action: Action, reward: float, next_state: GameState, end_reason: Optional[str] = None) -> None:
    """
    Adds a single step to the trajectory.

    Args:
        action (Action): The action taken.
        reward (float): The reward received.
        next_state (GameState): The resulting state.
        end_reason (Optional[str]): Reason for episode end, if applicable.
    """
    self.logger.debug(f"Adding step to trajectory for {self.agent_name}")
    if len(self._data["trajectory"]["states"]) == 0:
        self.logger.warning("The current action (id:{action.id}) is being added as the first step, but the initial state has not been recorded yet. Please call add_initial_state() at the beginning of the episode.")
    self._data["trajectory"]["actions"].append(action.as_dict)
    self._data["trajectory"]["rewards"].append(reward)
    self._data["trajectory"]["states"].append(next_state.as_dict)

    if end_reason:
        self._data["end_reason"] = end_reason

get_trajectory

Returns the current trajectory data.

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: The trajectory dictionary.

Source code in netsecgame/utils/trajectory_recorder.py

def get_trajectory(self) -> Dict[str, Any]:
    """
    Returns the current trajectory data.

    Returns:
        Dict[str, Any]: The trajectory dictionary.
    """
    return self._data

reset

Resets the trajectory data for a new episode.

Source code in netsecgame/utils/trajectory_recorder.py

def reset(self) -> None:
    """Resets the trajectory data for a new episode."""
    self.logger.debug(f"Resetting trajectory for {self.agent_name}")
    self._data = {
        "trajectory": {
            "states": [],
            "actions": [],
            "rewards": [],
        },
        "end_reason": None,
        "agent_role": self.agent_role,
        "agent_name": self.agent_name
    }

save_to_file

Saves the recorded trajectory to a JSONL file.

Parameters:

Name	Type	Description	Default
location	str	Directory to save the file. Defaults to "./logs/trajectories".	'./logs/trajectories'
filename	str	Name of the file to save. If None, defaults to "{datetime.now():%Y-%m-%d}{self.agent_name}{self.agent_role}".	None

Source code in netsecgame/utils/trajectory_recorder.py

def save_to_file(self, location: str = "./logs/trajectories", filename:str=None) -> None:
    """
    Saves the recorded trajectory to a JSONL file.

    Args:
        location (str): Directory to save the file. Defaults to "./logs/trajectories".
        filename (str): Name of the file to save. If None, defaults to "{datetime.now():%Y-%m-%d}_{self.agent_name}_{self.agent_role}".
    """
    if filename is None:
        filename = f"{datetime.now():%Y-%m-%d}_{self.agent_name}_{self.agent_role}"
    try:
        store_trajectories_to_jsonl(self._data, location, filename)
        self.logger.debug(f"Trajectory stored in {os.path.join(location, filename)}.jsonl")
    except Exception as e:
        self.logger.error(f"Failed to store trajectory: {e}")