Start RL in Colab with PyTorch and Gym

Check the PyTorch and Gym versions preinstalled in Colab

import torch
import gymnasium as gym

print(torch.__version__)
print(gym.__version__)

2.5.1+cu124
0.25.2

Run the sample code (Random CartPole)

• Simulate Cartpole with random control inputs
• Without effective control, the pole falls and the episode ends after only 1–2 steps.

def run_random_agent(env_name="CartPole-v1", max_episodes=2, max_steps=200):
    env = gym.make(env_name)
    for episode in range(max_episodes):
        state, info = env.reset()
        total_reward = 0
        for step in range(max_steps):
            action = env.action_space.sample()  # Random action
            next_state, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            state = next_state
            if terminated or truncated:
                print(f"Episode {episode+1} finished after {step+1} steps, total reward: {total_reward}")
                break
    env.close()

def tensor_example():
    x = torch.randn((3,4))
    print("PyTorch 텐서 예제:")
    print(x)
    print("텐서 크기:", x.shape)

if __name__ == "__main__":
    run_random_agent()
    tensor_example()

Episode 1 finished after 20 steps, total reward: 20.0
Episode 2 finished after 15 steps, total reward: 15.0
PyTorch 텐서 예제:
tensor([[-0.5245,  0.4341,  0.1545,  1.0109],
        [ 0.4827, -2.6886, -0.5962, -0.1738],
        [ 0.3569, -0.1071, -0.3224,  0.0959]])
텐서 크기: torch.Size([3, 4])