robots.py

import gymnasium as gym
import evogym.envs
from datasets import load_dataset
import numpy as np
from dataclasses import dataclass

@dataclass
class Robot:
    uid: str
    body: np.ndarray
    connections: np.ndarray
    reward: float
    env_name: str
    generated_by: str
    
    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            uid=data["uid"],
            body=np.array(data["body"], dtype=np.int64),
            connections=np.array(data["connections"], dtype=np.int64),
            reward=data["reward"],
            env_name=data["env_name"],
            generated_by=data["generated_by"],
        )

if __name__ == "__main__":
    dataset = load_dataset("EvoGym/robots")
    
    # Sample a random robot
    index = np.random.randint(0, len(dataset["train"]))
    robot_dict = dataset["train"][index]
    robot = Robot.from_dict(robot_dict)
    print(f"Sampled robot UID: {robot.uid}, at index {index}, and with env {robot.env_name}.")    
    
    # Create an environment with the robot and run a simulation
    env = gym.make(
        robot.env_name,
        body=robot.body,
        connections=robot.connections,
        render_mode='human'
    )
    env.reset()

    while True:
        action = env.action_space.sample()
        ob, reward, terminated, truncated, info = env.step(action)

        if terminated or truncated:
            break
        
    env.close()