Ngày Cộng đồng ML là ngày 9 tháng 11! Tham gia với chúng tôi để cập nhật từ TensorFlow, JAX, và nhiều hơn nữa Tìm hiểu thêm

Môi trường

Xem trên TensorFlow.org Xem nguồn trên GitHub Tải xuống sổ ghi chép

Giới thiệu

Mục tiêu của Học tăng cường (RL) là thiết kế các tác nhân học hỏi bằng cách tương tác với môi trường. Trong cài đặt RL tiêu chuẩn, tác nhân nhận được quan sát ở mỗi bước thời gian và chọn một hành động. Hành động được áp dụng cho môi trường và môi trường trả lại phần thưởng và một quan sát mới. Đại lý đào tạo một chính sách để lựa chọn các hành động để tối đa hóa tổng phần thưởng, còn được gọi là hoàn trả.

Trong TF-Agents, các môi trường có thể được triển khai bằng Python hoặc TensorFlow. Môi trường Python thường dễ triển khai, dễ hiểu và gỡ lỗi hơn, nhưng môi trường TensorFlow hiệu quả hơn và cho phép song song hóa tự nhiên. Quy trình công việc phổ biến nhất là triển khai môi trường bằng Python và sử dụng một trong các trình bao bọc của chúng tôi để tự động chuyển đổi nó thành TensorFlow.

Trước tiên, chúng ta hãy xem xét các môi trường Python. Môi trường TensorFlow tuân theo một API rất giống nhau.

Cài đặt

Nếu bạn chưa cài đặt tf-agent hoặc phòng tập thể dục, hãy chạy:

pip install "gym>=0.21.0"
pip install tf-agents
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import tensorflow as tf
import numpy as np

from tf_agents.environments import py_environment
from tf_agents.environments import tf_environment
from tf_agents.environments import tf_py_environment
from tf_agents.environments import utils
from tf_agents.specs import array_spec
from tf_agents.environments import wrappers
from tf_agents.environments import suite_gym
from tf_agents.trajectories import time_step as ts

Môi trường Python

Môi trường Python có một step(action) -> next_time_step phương pháp áp dụng một hành động đối với môi trường, và trả về thông tin sau về các bước tiếp theo:

  1. observation : Đây là một phần của tình trạng môi trường mà các đại lý có thể quan sát để lựa chọn hành động của mình ở bước tiếp theo.
  2. reward : Các đại lý đang học để tối đa hóa tổng của các phần thưởng qua nhiều bước.
  3. step_type : Tương tác với môi trường thường là một phần của một chuỗi / tập phim. ví dụ: nhiều nước đi trong một trò chơi cờ vua. step_type có thể là FIRST , MID hay LAST để cho biết bước thời gian này là bước đầu tiên, trung cấp hoặc cuối cùng trong một chuỗi.
  4. discount : Đây là một phao đại diện cho bao nhiêu trọng lượng phần thưởng đồng thời tiếp bước tương đối so với phần thưởng ở bước thời điểm hiện tại.

Đây là những nhóm lại thành một tuple tên TimeStep(step_type, reward, discount, observation) .

Giao diện rằng tất cả các môi trường Python phải thực hiện là trong environments/py_environment.PyEnvironment . Các phương pháp chính là:

class PyEnvironment(object):

  def reset(self):
    """Return initial_time_step."""
    self._current_time_step = self._reset()
    return self._current_time_step

  def step(self, action):
    """Apply action and return new time_step."""
    if self._current_time_step is None:
        return self.reset()
    self._current_time_step = self._step(action)
    return self._current_time_step

  def current_time_step(self):
    return self._current_time_step

  def time_step_spec(self):
    """Return time_step_spec."""

  @abc.abstractmethod
  def observation_spec(self):
    """Return observation_spec."""

  @abc.abstractmethod
  def action_spec(self):
    """Return action_spec."""

  @abc.abstractmethod
  def _reset(self):
    """Return initial_time_step."""

  @abc.abstractmethod
  def _step(self, action):
    """Apply action and return new time_step."""

Ngoài các step() phương pháp, môi trường cũng cung cấp một reset() phương pháp bắt đầu một chuỗi mới và cung cấp một đầu TimeStep . Nó không phải là cần thiết để gọi reset phương pháp một cách rõ ràng. Chúng tôi giả định rằng các môi trường được đặt lại tự động, khi chúng đi đến cuối một tập hoặc khi bước () được gọi lần đầu tiên.

Lưu ý rằng lớp con không thực hiện step() hoặc reset() trực tiếp. Họ thay vì ghi đè _step()_reset() phương pháp. Các bước thời gian trở về từ những phương pháp này sẽ được lưu trữ và tiếp xúc qua current_time_step() .

Các observation_specaction_spec phương pháp trả về một tổ của (Bounded)ArraySpecs mô tả tên, hình dáng, kiểu dữ liệu và phạm vi của các quan sát và hành động tương ứng.

Trong TF-Agents, chúng tôi liên tục đề cập đến các tổ được định nghĩa là bất kỳ cây nào giống như cấu trúc bao gồm danh sách, bộ giá trị, bộ mã được đặt tên hoặc từ điển. Chúng có thể được tạo tùy ý để duy trì cấu trúc của các quan sát và hành động. Chúng tôi nhận thấy điều này rất hữu ích cho các môi trường phức tạp hơn, nơi bạn có nhiều quan sát và hành động.

Sử dụng Môi trường Tiêu chuẩn

TF Đại lý đã được xây dựng trong giấy gói cho nhiều môi trường tiêu chuẩn như các OpenAI phòng tập thể dục, DeepMind kiểm soát và Atari, để họ làm theo chúng tôi py_environment.PyEnvironment giao diện. Các môi trường được bao bọc này có thể được tải dễ dàng bằng cách sử dụng các bộ môi trường của chúng tôi. Hãy tải môi trường CartPole từ phòng tập thể dục OpenAI và xem xét action và time_step_spec.

environment = suite_gym.load('CartPole-v0')
print('action_spec:', environment.action_spec())
print('time_step_spec.observation:', environment.time_step_spec().observation)
print('time_step_spec.step_type:', environment.time_step_spec().step_type)
print('time_step_spec.discount:', environment.time_step_spec().discount)
print('time_step_spec.reward:', environment.time_step_spec().reward)
action_spec: BoundedArraySpec(shape=(), dtype=dtype('int64'), name='action', minimum=0, maximum=1)
time_step_spec.observation: BoundedArraySpec(shape=(4,), dtype=dtype('float32'), name='observation', minimum=[-4.8000002e+00 -3.4028235e+38 -4.1887903e-01 -3.4028235e+38], maximum=[4.8000002e+00 3.4028235e+38 4.1887903e-01 3.4028235e+38])
time_step_spec.step_type: ArraySpec(shape=(), dtype=dtype('int32'), name='step_type')
time_step_spec.discount: BoundedArraySpec(shape=(), dtype=dtype('float32'), name='discount', minimum=0.0, maximum=1.0)
time_step_spec.reward: ArraySpec(shape=(), dtype=dtype('float32'), name='reward')

Vì vậy, chúng ta thấy rằng môi trường hy vọng hành động của loại int64 trong [0, 1] và lợi nhuận TimeSteps nơi các quan sát là một float32 vector có độ dài 4 và yếu tố giảm giá là một float32 trong [0.0, 1.0]. Bây giờ, chúng ta hãy cố gắng để có một hành động cố định (1,) cho cả một tập phim.

action = np.array(1, dtype=np.int32)
time_step = environment.reset()
print(time_step)
while not time_step.is_last():
  time_step = environment.step(action)
  print(time_step)
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.0138565 , -0.03582913,  0.04861612, -0.03755046], dtype=float32),
 'reward': array(0., dtype=float32),
 'step_type': array(0, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.01313992,  0.15856317,  0.0478651 , -0.3145069 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.01631118,  0.35297176,  0.04157497, -0.5917188 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.02337062,  0.54748774,  0.02974059, -0.87102115], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.03432037,  0.74219286,  0.01232017, -1.1542072 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.04916423,  0.93715197, -0.01076398, -1.4430016 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.06790727,  1.1324048 , -0.03962401, -1.7390285 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.09055536,  1.327955  , -0.07440457, -2.04377   ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.11711447,  1.523758  , -0.11527998, -2.3585167 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([ 0.14758962,  1.7197047 , -0.16245031, -2.6843033 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([ 0.18198372,  1.9156038 , -0.21613638, -3.0218334 ], dtype=float32),
 'reward': array(1., dtype=float32),
 'step_type': array(2, dtype=int32)})

Tạo môi trường Python của riêng bạn

Đối với nhiều khách hàng, trường hợp sử dụng phổ biến là áp dụng một trong các tác nhân tiêu chuẩn (xem các đại lý /) trong TF-Agents cho vấn đề của họ. Để làm được điều này, họ phải định hình vấn đề của mình như một môi trường. Vì vậy, chúng ta hãy xem cách triển khai một môi trường trong Python.

Giả sử chúng tôi muốn đào tạo một đặc vụ để chơi trò chơi bài (lấy cảm hứng từ Black Jack) sau:

  1. Trò chơi được chơi bằng cách sử dụng một bộ bài vô hạn được đánh số từ 1 ... 10.
  2. Tại mỗi lượt, nhân viên có thể làm 2 việc: nhận một thẻ ngẫu nhiên mới, hoặc dừng vòng hiện tại.
  3. Mục tiêu là nhận được tổng số thẻ của bạn càng gần 21 càng tốt vào cuối vòng mà không vượt quá.

Môi trường đại diện cho trò chơi có thể trông như thế này:

  1. Hành động: Chúng tôi có 2 hành động. Hành động 0: nhận thẻ mới và Hành động 1: chấm dứt vòng hiện tại.
  2. Quan sát: Tổng của các thẻ trong vòng hiện tại.
  3. Phần thưởng: Mục tiêu là đạt được càng gần 21 càng tốt mà không cần vượt quá, vì vậy chúng tôi có thể đạt được điều này bằng cách sử dụng phần thưởng sau vào cuối vòng: sum_of_cards - 21 nếu sum_of_cards <= 21, khác -21
class CardGameEnv(py_environment.PyEnvironment):

  def __init__(self):
    self._action_spec = array_spec.BoundedArraySpec(
        shape=(), dtype=np.int32, minimum=0, maximum=1, name='action')
    self._observation_spec = array_spec.BoundedArraySpec(
        shape=(1,), dtype=np.int32, minimum=0, name='observation')
    self._state = 0
    self._episode_ended = False

  def action_spec(self):
    return self._action_spec

  def observation_spec(self):
    return self._observation_spec

  def _reset(self):
    self._state = 0
    self._episode_ended = False
    return ts.restart(np.array([self._state], dtype=np.int32))

  def _step(self, action):

    if self._episode_ended:
      # The last action ended the episode. Ignore the current action and start
      # a new episode.
      return self.reset()

    # Make sure episodes don't go on forever.
    if action == 1:
      self._episode_ended = True
    elif action == 0:
      new_card = np.random.randint(1, 11)
      self._state += new_card
    else:
      raise ValueError('`action` should be 0 or 1.')

    if self._episode_ended or self._state >= 21:
      reward = self._state - 21 if self._state <= 21 else -21
      return ts.termination(np.array([self._state], dtype=np.int32), reward)
    else:
      return ts.transition(
          np.array([self._state], dtype=np.int32), reward=0.0, discount=1.0)

Hãy đảm bảo rằng chúng tôi đã làm mọi thứ xác định đúng môi trường ở trên. Khi tạo môi trường của riêng bạn, bạn phải đảm bảo rằng các quan sát và bước thời gian được tạo ra tuân theo các hình dạng và kiểu chính xác như được xác định trong thông số kỹ thuật của bạn. Chúng được sử dụng để tạo biểu đồ TensorFlow và như vậy có thể tạo ra các vấn đề khó gỡ lỗi nếu chúng tôi làm sai.

Để xác thực môi trường của chúng tôi, chúng tôi sẽ sử dụng một chính sách ngẫu nhiên để tạo các hành động và chúng tôi sẽ lặp lại hơn 5 tập để đảm bảo mọi thứ đang hoạt động như dự định. Lỗi sẽ xảy ra nếu chúng ta nhận được time_step không tuân theo các thông số kỹ thuật của môi trường.

environment = CardGameEnv()
utils.validate_py_environment(environment, episodes=5)

Bây giờ chúng ta biết môi trường đang hoạt động như dự định, hãy chạy môi trường này bằng một chính sách cố định: yêu cầu 3 thẻ và sau đó kết thúc vòng chơi.

get_new_card_action = np.array(0, dtype=np.int32)
end_round_action = np.array(1, dtype=np.int32)

environment = CardGameEnv()
time_step = environment.reset()
print(time_step)
cumulative_reward = time_step.reward

for _ in range(3):
  time_step = environment.step(get_new_card_action)
  print(time_step)
  cumulative_reward += time_step.reward

time_step = environment.step(end_round_action)
print(time_step)
cumulative_reward += time_step.reward
print('Final Reward = ', cumulative_reward)
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([0], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(0, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([9], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(1., dtype=float32),
 'observation': array([12], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(1, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([21], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(2, dtype=int32)})
TimeStep(
{'discount': array(0., dtype=float32),
 'observation': array([21], dtype=int32),
 'reward': array(0., dtype=float32),
 'step_type': array(2, dtype=int32)})
Final Reward =  0.0

Môi trường Wrappers

Một trình bao bọc môi trường lấy một môi trường Python và trả về một phiên bản đã sửa đổi của môi trường. Cả hai môi trường ban đầu và môi trường biến đổi là trường hợp của py_environment.PyEnvironment , và nhiều giấy gói có thể được xích lại với nhau.

Một số giấy gói thông thường có thể được tìm thấy trong environments/wrappers.py . Ví dụ:

  1. ActionDiscretizeWrapper : Chuyển đổi một không gian hoạt động liên tục đến một không gian hành động rời rạc.
  2. RunStats : Captures chạy thống kê của môi trường như số lượng các bước thực hiện, số lượng tập phim hoàn thành, vv
  3. TimeLimit : Kết thúc các tập phim sau khi một số cố định các bước.

Ví dụ 1: Hành động riêng biệt với trình bao bọc

InvertedPendulum là một môi trường PyBullet chấp nhận hành động liên tục trong khoảng [-2, 2] . Nếu chúng ta muốn đào tạo một tác nhân hành động rời rạc như DQN trên môi trường này, chúng ta phải rời rạc hóa (lượng tử hóa) không gian hành động. Đây là chính xác những gì ActionDiscretizeWrapper làm. Hãy so sánh các action_spec trước và sau khi gói:

env = suite_gym.load('Pendulum-v1')
print('Action Spec:', env.action_spec())

discrete_action_env = wrappers.ActionDiscretizeWrapper(env, num_actions=5)
print('Discretized Action Spec:', discrete_action_env.action_spec())
Action Spec: BoundedArraySpec(shape=(1,), dtype=dtype('float32'), name='action', minimum=-2.0, maximum=2.0)
Discretized Action Spec: BoundedArraySpec(shape=(), dtype=dtype('int32'), name='action', minimum=0, maximum=4)

Các gói discrete_action_env là một thể hiện của py_environment.PyEnvironment và có thể bị đối xử như một môi trường Python thường xuyên.

Môi trường TensorFlow

Giao diện cho các môi trường TF được định nghĩa trong environments/tf_environment.TFEnvironment và ngoại hình rất giống với môi trường Python. Môi trường TF khác với môi trường Python theo một số cách:

  • Chúng tạo ra các đối tượng tensor thay vì mảng
  • Môi trường TF thêm một thứ nguyên lô vào lực căng được tạo ra khi so sánh với thông số kỹ thuật.

Việc chuyển đổi môi trường Python thành TFEnvs cho phép tensorflow song song hóa các hoạt động. Ví dụ, người ta có thể xác định một collect_experience_op thu thập dữ liệu từ môi trường, bổ sung một replay_buffer , và một train_op mà đọc từ replay_buffer và huấn luyện các đại lý, và chạy chúng song song một cách tự nhiên trong TensorFlow.

class TFEnvironment(object):

  def time_step_spec(self):
    """Describes the `TimeStep` tensors returned by `step()`."""

  def observation_spec(self):
    """Defines the `TensorSpec` of observations provided by the environment."""

  def action_spec(self):
    """Describes the TensorSpecs of the action expected by `step(action)`."""

  def reset(self):
    """Returns the current `TimeStep` after resetting the Environment."""
    return self._reset()

  def current_time_step(self):
    """Returns the current `TimeStep`."""
    return self._current_time_step()

  def step(self, action):
    """Applies the action and returns the new `TimeStep`."""
    return self._step(action)

  @abc.abstractmethod
  def _reset(self):
    """Returns the current `TimeStep` after resetting the Environment."""

  @abc.abstractmethod
  def _current_time_step(self):
    """Returns the current `TimeStep`."""

  @abc.abstractmethod
  def _step(self, action):
    """Applies the action and returns the new `TimeStep`."""

Các current_time_step() phương thức trả về các time_step hiện tại và khởi môi trường nếu cần thiết.

Các reset() phương pháp lực thiết lập lại trong môi trường và lợi nhuận các CURRENT_STEP.

Nếu action không phụ thuộc vào trước time_step một tf.control_dependency là cần thiết trong Graph chế độ.

Còn bây giờ, chúng ta hãy nhìn vào cách TFEnvironments được tạo ra.

Tạo môi trường TensorFlow của riêng bạn

Điều này phức tạp hơn so với việc tạo môi trường trong Python, vì vậy chúng tôi sẽ không đề cập đến nó trong chuyên mục này. Một ví dụ có sẵn ở đây . Các trường hợp sử dụng phổ biến hơn là thực hiện môi trường của bạn bằng Python và bọc nó trong TensorFlow sử dụng của chúng tôi TFPyEnvironment wrapper (xem dưới đây).

Gói một môi trường Python trong TensorFlow

Chúng ta có thể dễ dàng quấn bất kỳ môi trường Python vào một môi trường TensorFlow sử dụng TFPyEnvironment wrapper.

env = suite_gym.load('CartPole-v0')
tf_env = tf_py_environment.TFPyEnvironment(env)

print(isinstance(tf_env, tf_environment.TFEnvironment))
print("TimeStep Specs:", tf_env.time_step_spec())
print("Action Specs:", tf_env.action_spec())
True
TimeStep Specs: TimeStep(
{'discount': BoundedTensorSpec(shape=(), dtype=tf.float32, name='discount', minimum=array(0., dtype=float32), maximum=array(1., dtype=float32)),
 'observation': BoundedTensorSpec(shape=(4,), dtype=tf.float32, name='observation', minimum=array([-4.8000002e+00, -3.4028235e+38, -4.1887903e-01, -3.4028235e+38],
      dtype=float32), maximum=array([4.8000002e+00, 3.4028235e+38, 4.1887903e-01, 3.4028235e+38],
      dtype=float32)),
 'reward': TensorSpec(shape=(), dtype=tf.float32, name='reward'),
 'step_type': TensorSpec(shape=(), dtype=tf.int32, name='step_type')})
Action Specs: BoundedTensorSpec(shape=(), dtype=tf.int64, name='action', minimum=array(0), maximum=array(1))

Lưu ý các thông số kỹ thuật hiện nay có các loại: (Bounded)TensorSpec .

Ví dụ sử dụng

Ví dụ đơn giản

env = suite_gym.load('CartPole-v0')

tf_env = tf_py_environment.TFPyEnvironment(env)
# reset() creates the initial time_step after resetting the environment.
time_step = tf_env.reset()
num_steps = 3
transitions = []
reward = 0
for i in range(num_steps):
  action = tf.constant([i % 2])
  # applies the action and returns the new TimeStep.
  next_time_step = tf_env.step(action)
  transitions.append([time_step, action, next_time_step])
  reward += next_time_step.reward
  time_step = next_time_step

np_transitions = tf.nest.map_structure(lambda x: x.numpy(), transitions)
print('\n'.join(map(str, np_transitions)))
print('Total reward:', reward.numpy())
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.0078796 , -0.04736348, -0.04966116,  0.04563603]],
      dtype=float32),
 'reward': array([0.], dtype=float32),
 'step_type': array([0], dtype=int32)}), array([0], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.00882687, -0.24173944, -0.04874843,  0.32224613]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.00882687, -0.24173944, -0.04874843,  0.32224613]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)}), array([1], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01366166, -0.04595843, -0.04230351,  0.01459712]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
[TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01366166, -0.04595843, -0.04230351,  0.01459712]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)}), array([0], dtype=int32), TimeStep(
{'discount': array([1.], dtype=float32),
 'observation': array([[-0.01458083, -0.24044897, -0.04201157,  0.2936384 ]],
      dtype=float32),
 'reward': array([1.], dtype=float32),
 'step_type': array([1], dtype=int32)})]
Total reward: [3.]

Toàn bộ tập

env = suite_gym.load('CartPole-v0')
tf_env = tf_py_environment.TFPyEnvironment(env)

time_step = tf_env.reset()
rewards = []
steps = []
num_episodes = 5

for _ in range(num_episodes):
  episode_reward = 0
  episode_steps = 0
  while not time_step.is_last():
    action = tf.random.uniform([1], 0, 2, dtype=tf.int32)
    time_step = tf_env.step(action)
    episode_steps += 1
    episode_reward += time_step.reward.numpy()
  rewards.append(episode_reward)
  steps.append(episode_steps)
  time_step = tf_env.reset()

num_steps = np.sum(steps)
avg_length = np.mean(steps)
avg_reward = np.mean(rewards)

print('num_episodes:', num_episodes, 'num_steps:', num_steps)
print('avg_length', avg_length, 'avg_reward:', avg_reward)
num_episodes: 5 num_steps: 131
avg_length 26.2 avg_reward: 26.2